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

95 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-03-30 03:34 -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.deblender import SourceDeblendTask 

32from lsst.meas.extensions.scarlet import ScarletDeblendTask 

33from lsst.obs.base import ExposureIdInfo 

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 def setDefaults(self): 

82 super().setDefaults() 

83 self.singleBandDeblend.propagateAllPeaks = True 

84 

85 

86class DeblendCoaddSourcesSingleConfig(PipelineTaskConfig, 

87 pipelineConnections=DeblendCoaddSourceSingleConnections): 

88 singleBandDeblend = ConfigurableField( 

89 target=SourceDeblendTask, 

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

91 ) 

92 

93 

94class DeblendCoaddSourcesMultiConnections(PipelineTaskConnections, 

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

96 defaultTemplates=deblendBaseTemplates): 

97 inputSchema = cT.InitInput( 

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

99 name="{inputCoaddName}Coadd_mergeDet_schema", 

100 storageClass="SourceCatalog" 

101 ) 

102 peakSchema = cT.InitInput( 

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

104 name="{inputCoaddName}Coadd_peak_schema", 

105 storageClass="PeakCatalog" 

106 ) 

107 mergedDetections = cT.Input( 

108 doc="Detection catalog merged across bands", 

109 name="{inputCoaddName}Coadd_mergeDet", 

110 storageClass="SourceCatalog", 

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

112 ) 

113 coadds = cT.Input( 

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

115 name="{inputCoaddName}Coadd_calexp", 

116 storageClass="ExposureF", 

117 multiple=True, 

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

119 ) 

120 outputSchema = cT.InitOutput( 

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

122 name="{outputCoaddName}Coadd_deblendedFlux_schema", 

123 storageClass="SourceCatalog" 

124 ) 

125 fluxCatalogs = cT.Output( 

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

127 name="{outputCoaddName}Coadd_deblendedFlux", 

128 storageClass="SourceCatalog", 

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

130 multiple=True 

131 ) 

132 templateCatalogs = cT.Output( 

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

134 name="{outputCoaddName}Coadd_deblendedModel", 

135 storageClass="SourceCatalog", 

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

137 multiple=True 

138 ) 

139 deblendedCatalog = cT.Output( 

140 doc="Catalogs produced by multiband deblending", 

141 name="{outputCoaddName}Coadd_deblendedCatalog", 

142 storageClass="SourceCatalog", 

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

144 ) 

145 scarletModelData = cT.Output( 

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

147 name="{outputCoaddName}Coadd_scarletModelData", 

148 storageClass="ScarletModelData", 

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

150 ) 

151 

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

153 super().__init__(config=config) 

154 # Remove unused connections. 

155 # TODO: deprecate once RFC-860 passes. 

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

157 

158 

159class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig, 

160 pipelineConnections=DeblendCoaddSourcesMultiConnections): 

161 multibandDeblend = ConfigurableField( 

162 target=ScarletDeblendTask, 

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

164 ) 

165 

166 

167class DeblendCoaddSourcesBaseTask(PipelineTask): 

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

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

170 schema = initInputs["inputSchema"].schema 

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

172 self.schemaMapper = afwTable.SchemaMapper(schema) 

173 self.schemaMapper.addMinimalSchema(schema) 

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

175 

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

177 inputs = butlerQC.get(inputRefs) 

178 inputs["idFactory"] = ExposureIdInfo.fromDataId( 

179 butlerQC.quantum.dataId, 

180 "tract_patch" 

181 ).makeSourceIdFactory() 

182 outputs = self.run(**inputs) 

183 butlerQC.put(outputs, outputRefs) 

184 

185 def _makeSourceCatalog(self, mergedDetections, idFactory): 

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

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

188 # with the largest id already in the catalog. 

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

190 idFactory.notify(maxId) 

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

192 sources = afwTable.SourceCatalog(table) 

193 sources.extend(mergedDetections, self.schemaMapper) 

194 return sources 

195 

196 

197class DeblendCoaddSourcesSingleTask(DeblendCoaddSourcesBaseTask): 

198 ConfigClass = DeblendCoaddSourcesSingleConfig 

199 _DefaultName = "deblendCoaddSourcesSingle" 

200 

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

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

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

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

205 

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

207 sources = self._makeSourceCatalog(mergedDetections, idFactory) 

208 self.singleBandDeblend.run(coadd, sources) 

209 if not sources.isContiguous(): 

210 sources = sources.copy(deep=True) 

211 return Struct(measureCatalog=sources) 

212 

213 

214class DeblendCoaddSourcesMultiTask(DeblendCoaddSourcesBaseTask): 

215 ConfigClass = DeblendCoaddSourcesMultiConfig 

216 _DefaultName = "deblendCoaddSourcesMulti" 

217 

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

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

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

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

222 

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

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

225 # all input lists to match this band order. 

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

227 bandOrder.sort() 

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

229 inputs = butlerQC.get(inputRefs) 

230 exposureIdInfo = ExposureIdInfo.fromDataId(butlerQC.quantum.dataId, "tract_patch") 

231 inputs["idFactory"] = exposureIdInfo.makeSourceIdFactory() 

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

233 outputs = self.run(**inputs) 

234 butlerQC.put(outputs, outputRefs) 

235 

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

237 sources = self._makeSourceCatalog(mergedDetections, idFactory) 

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

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

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

241 return retStruct