Coverage for python/lsst/pipe/tasks/deblendCoaddSourcesPipeline.py: 51%

96 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-06-07 02:55 -0700

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__ = ["DeblendCoaddSourcesSingleConfig", "DeblendCoaddSourcesSingleTask", 

23 "DeblendCoaddSourcesMultiConfig", "DeblendCoaddSourcesMultiTask"] 

24 

25import numpy as np 

26 

27from lsst.pipe.base import (Struct, PipelineTask, PipelineTaskConfig, PipelineTaskConnections) 

28import lsst.pipe.base.connectionTypes as cT 

29 

30from lsst.pex.config import ConfigurableField 

31from lsst.meas.base import SkyMapIdGeneratorConfig 

32from lsst.meas.deblender import SourceDeblendTask 

33from lsst.meas.extensions.scarlet import ScarletDeblendTask 

34 

35import lsst.afw.image as afwImage 

36import lsst.afw.table as afwTable 

37 

38from .makeWarp import reorderRefs 

39 

40 

41deblendBaseTemplates = {"inputCoaddName": "deep", "outputCoaddName": "deep"} 

42 

43 

44class DeblendCoaddSourceSingleConnections(PipelineTaskConnections, 

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

46 defaultTemplates=deblendBaseTemplates): 

47 inputSchema = cT.InitInput( 

48 doc="Input schema to use in the deblend catalog", 

49 name="{inputCoaddName}Coadd_mergeDet_schema", 

50 storageClass="SourceCatalog" 

51 ) 

52 peakSchema = cT.InitInput( 

53 doc="Schema of the footprint peak catalogs", 

54 name="{inputCoaddName}Coadd_peak_schema", 

55 storageClass="PeakCatalog" 

56 ) 

57 mergedDetections = cT.Input( 

58 doc="Detection catalog merged across bands", 

59 name="{inputCoaddName}Coadd_mergeDet", 

60 storageClass="SourceCatalog", 

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

62 ) 

63 coadd = cT.Input( 

64 doc="Exposure on which to run deblending", 

65 name="{inputCoaddName}Coadd_calexp", 

66 storageClass="ExposureF", 

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

68 ) 

69 measureCatalog = cT.Output( 

70 doc="The output measurement catalog of deblended sources", 

71 name="{outputCoaddName}Coadd_deblendedFlux", 

72 storageClass="SourceCatalog", 

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

74 ) 

75 outputSchema = cT.InitOutput( 

76 doc="Output of the schema used in deblending task", 

77 name="{outputCoaddName}Coadd_deblendedFlux_schema", 

78 storageClass="SourceCatalog" 

79 ) 

80 

81 

82class DeblendCoaddSourcesSingleConfig(PipelineTaskConfig, 

83 pipelineConnections=DeblendCoaddSourceSingleConnections): 

84 singleBandDeblend = ConfigurableField( 

85 target=SourceDeblendTask, 

86 doc="Task to deblend an image in one band" 

87 ) 

88 idGenerator = SkyMapIdGeneratorConfig.make_field() 

89 

90 def setDefaults(self): 

91 super().setDefaults() 

92 self.singleBandDeblend.propagateAllPeaks = True 

93 

94 

95class DeblendCoaddSourcesMultiConnections(PipelineTaskConnections, 

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

97 defaultTemplates=deblendBaseTemplates): 

98 inputSchema = cT.InitInput( 

99 doc="Input schema to use in the deblend catalog", 

100 name="{inputCoaddName}Coadd_mergeDet_schema", 

101 storageClass="SourceCatalog" 

102 ) 

103 peakSchema = cT.InitInput( 

104 doc="Schema of the footprint peak catalogs", 

105 name="{inputCoaddName}Coadd_peak_schema", 

106 storageClass="PeakCatalog" 

107 ) 

108 mergedDetections = cT.Input( 

109 doc="Detection catalog merged across bands", 

110 name="{inputCoaddName}Coadd_mergeDet", 

111 storageClass="SourceCatalog", 

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

113 ) 

114 coadds = cT.Input( 

115 doc="Exposure on which to run deblending", 

116 name="{inputCoaddName}Coadd_calexp", 

117 storageClass="ExposureF", 

118 multiple=True, 

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

120 ) 

121 outputSchema = cT.InitOutput( 

122 doc="Output of the schema used in deblending task", 

123 name="{outputCoaddName}Coadd_deblendedFlux_schema", 

124 storageClass="SourceCatalog" 

125 ) 

126 fluxCatalogs = cT.Output( 

127 doc="Flux weighted catalogs produced by multiband deblending", 

128 name="{outputCoaddName}Coadd_deblendedFlux", 

129 storageClass="SourceCatalog", 

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

131 multiple=True 

132 ) 

133 templateCatalogs = cT.Output( 

134 doc="Template catalogs produced by multiband deblending", 

135 name="{outputCoaddName}Coadd_deblendedModel", 

136 storageClass="SourceCatalog", 

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

138 multiple=True 

139 ) 

140 deblendedCatalog = cT.Output( 

141 doc="Catalogs produced by multiband deblending", 

142 name="{outputCoaddName}Coadd_deblendedCatalog", 

143 storageClass="SourceCatalog", 

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

145 ) 

146 scarletModelData = cT.Output( 

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

148 name="{outputCoaddName}Coadd_scarletModelData", 

149 storageClass="ScarletModelData", 

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

151 ) 

152 

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

154 super().__init__(config=config) 

155 # Remove unused connections. 

156 # TODO: deprecate once RFC-860 passes. 

157 self.outputs -= set(("fluxCatalogs", "templateCatalogs")) 

158 

159 

160class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig, 

161 pipelineConnections=DeblendCoaddSourcesMultiConnections): 

162 multibandDeblend = ConfigurableField( 

163 target=ScarletDeblendTask, 

164 doc="Task to deblend an images in multiple bands" 

165 ) 

166 idGenerator = SkyMapIdGeneratorConfig.make_field() 

167 

168 

169class DeblendCoaddSourcesBaseTask(PipelineTask): 

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

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

172 schema = initInputs["inputSchema"].schema 

173 self.peakSchema = initInputs["peakSchema"].schema 

174 self.schemaMapper = afwTable.SchemaMapper(schema) 

175 self.schemaMapper.addMinimalSchema(schema) 

176 self.schema = self.schemaMapper.getOutputSchema() 

177 

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

179 inputs = butlerQC.get(inputRefs) 

180 inputs["idFactory"] = self.config.idGenerator.apply(butlerQC.quantum.dataId).make_table_id_factory() 

181 outputs = self.run(**inputs) 

182 butlerQC.put(outputs, outputRefs) 

183 

184 def _makeSourceCatalog(self, mergedDetections, idFactory): 

185 # There may be gaps in the mergeDet catalog, which will cause the 

186 # source ids to be inconsistent. So we update the id factory 

187 # with the largest id already in the catalog. 

188 maxId = np.max(mergedDetections["id"]) 

189 idFactory.notify(maxId) 

190 table = afwTable.SourceTable.make(self.schema, idFactory) 

191 sources = afwTable.SourceCatalog(table) 

192 sources.extend(mergedDetections, self.schemaMapper) 

193 return sources 

194 

195 

196class DeblendCoaddSourcesSingleTask(DeblendCoaddSourcesBaseTask): 

197 ConfigClass = DeblendCoaddSourcesSingleConfig 

198 _DefaultName = "deblendCoaddSourcesSingle" 

199 

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

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

202 self.makeSubtask("singleBandDeblend", schema=self.schema, peakSchema=self.peakSchema) 

203 self.outputSchema = afwTable.SourceCatalog(self.schema) 

204 

205 def run(self, coadd, mergedDetections, idFactory): 

206 sources = self._makeSourceCatalog(mergedDetections, idFactory) 

207 self.singleBandDeblend.run(coadd, sources) 

208 if not sources.isContiguous(): 

209 sources = sources.copy(deep=True) 

210 return Struct(measureCatalog=sources) 

211 

212 

213class DeblendCoaddSourcesMultiTask(DeblendCoaddSourcesBaseTask): 

214 ConfigClass = DeblendCoaddSourcesMultiConfig 

215 _DefaultName = "deblendCoaddSourcesMulti" 

216 

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

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

219 self.makeSubtask("multibandDeblend", schema=self.schema, peakSchema=self.peakSchema) 

220 self.outputSchema = afwTable.SourceCatalog(self.schema) 

221 

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

223 # Obtain the list of bands, sort them (alphabetically), then reorder 

224 # all input lists to match this band order. 

225 bandOrder = [dRef.dataId["band"] for dRef in inputRefs.coadds] 

226 bandOrder.sort() 

227 inputRefs = reorderRefs(inputRefs, bandOrder, dataIdKey="band") 

228 inputs = butlerQC.get(inputRefs) 

229 inputs["idFactory"] = self.config.idGenerator.apply(butlerQC.quantum.dataId).make_table_id_factory() 

230 inputs["filters"] = [dRef.dataId["band"] for dRef in inputRefs.coadds] 

231 outputs = self.run(**inputs) 

232 butlerQC.put(outputs, outputRefs) 

233 

234 def run(self, coadds, filters, mergedDetections, idFactory): 

235 sources = self._makeSourceCatalog(mergedDetections, idFactory) 

236 multiExposure = afwImage.MultibandExposure.fromExposures(filters, coadds) 

237 catalog, modelData = self.multibandDeblend.run(multiExposure, sources) 

238 retStruct = Struct(deblendedCatalog=catalog, scarletModelData=modelData) 

239 return retStruct