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

Shortcuts on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

97 statements  

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 

22import numpy as np 

23 

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

25import lsst.pipe.base.connectionTypes as cT 

26 

27from lsst.pex.config import ConfigurableField 

28from lsst.meas.deblender import SourceDeblendTask 

29from lsst.meas.extensions.scarlet import ScarletDeblendTask 

30from lsst.obs.base import ExposureIdInfo 

31 

32import lsst.afw.image as afwImage 

33import lsst.afw.table as afwTable 

34 

35from .makeCoaddTempExp import reorderRefs 

36 

37__all__ = ("DeblendCoaddSourcesSingleConfig", "DeblendCoaddSourcesSingleTask", 

38 "DeblendCoaddSourcesMultiConfig", "DeblendCoaddSourcesMultiTask") 

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 

140 

141class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig, 

142 pipelineConnections=DeblendCoaddSourcesMultiConnections): 

143 multibandDeblend = ConfigurableField( 

144 target=ScarletDeblendTask, 

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

146 ) 

147 

148 

149class DeblendCoaddSourcesBaseTask(PipelineTask): 

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

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

152 schema = initInputs["inputSchema"].schema 

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

154 self.schemaMapper = afwTable.SchemaMapper(schema) 

155 self.schemaMapper.addMinimalSchema(schema) 

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

157 

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

159 inputs = butlerQC.get(inputRefs) 

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

161 butlerQC.quantum.dataId, 

162 "tract_patch" 

163 ).makeSourceIdFactory() 

164 outputs = self.run(**inputs) 

165 butlerQC.put(outputs, outputRefs) 

166 

167 def _makeSourceCatalog(self, mergedDetections, idFactory): 

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

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

170 # with the largest id already in the catalog. 

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

172 idFactory.notify(maxId) 

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

174 sources = afwTable.SourceCatalog(table) 

175 sources.extend(mergedDetections, self.schemaMapper) 

176 return sources 

177 

178 

179class DeblendCoaddSourcesSingleTask(DeblendCoaddSourcesBaseTask): 

180 ConfigClass = DeblendCoaddSourcesSingleConfig 

181 _DefaultName = "deblendCoaddSourcesSingle" 

182 

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

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

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

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

187 

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

189 sources = self._makeSourceCatalog(mergedDetections, idFactory) 

190 self.singleBandDeblend.run(coadd, sources) 

191 if not sources.isContiguous(): 

192 sources = sources.copy(deep=True) 

193 return Struct(measureCatalog=sources) 

194 

195 

196class DeblendCoaddSourcesMultiTask(DeblendCoaddSourcesBaseTask): 

197 ConfigClass = DeblendCoaddSourcesMultiConfig 

198 _DefaultName = "deblendCoaddSourcesMulti" 

199 

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

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

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

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

204 

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

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

207 # all input lists to match this band order. 

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

209 bandOrder.sort() 

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

211 inputs = butlerQC.get(inputRefs) 

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

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

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

215 outputs = self.run(**inputs) 

216 for outRef in outputRefs.templateCatalogs: 

217 band = outRef.dataId['band'] 

218 if (catalog := outputs.templateCatalogs.get(band)) is not None: 

219 butlerQC.put(catalog, outRef) 

220 

221 for outRef in outputRefs.fluxCatalogs: 

222 band = outRef.dataId['band'] 

223 if (catalog := outputs.fluxCatalogs.get(band)) is not None: 

224 butlerQC.put(catalog, outRef) 

225 

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

227 sources = self._makeSourceCatalog(mergedDetections, idFactory) 

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

229 templateCatalogs, fluxCatalogs = self.multibandDeblend.run(multiExposure, sources) 

230 retStruct = Struct(templateCatalogs=templateCatalogs, fluxCatalogs=fluxCatalogs) 

231 return retStruct