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

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__ = ["DeblendCoaddSourcesMultiConfig", "DeblendCoaddSourcesMultiTask"] 

23 

24import numpy as np 

25 

26from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections 

27import lsst.pipe.base.connectionTypes as cT 

28 

29from lsst.pex.config import ConfigurableField, Field 

30from lsst.meas.base import SkyMapIdGeneratorConfig 

31from lsst.meas.extensions.scarlet import ScarletDeblendTask 

32 

33import lsst.afw.image as afwImage 

34import lsst.afw.table as afwTable 

35 

36from .coaddBase import reorderRefs 

37 

38 

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

40 

41 

42class DeblendCoaddSourcesMultiConnections(PipelineTaskConnections, 

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

44 defaultTemplates=deblendBaseTemplates): 

45 inputSchema = cT.InitInput( 

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

47 name="{inputCoaddName}Coadd_mergeDet_schema", 

48 storageClass="SourceCatalog" 

49 ) 

50 peakSchema = cT.InitInput( 

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

52 name="{inputCoaddName}Coadd_peak_schema", 

53 storageClass="PeakCatalog" 

54 ) 

55 mergedDetections = cT.Input( 

56 doc="Detection catalog merged across bands", 

57 name="{inputCoaddName}Coadd_mergeDet", 

58 storageClass="SourceCatalog", 

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

60 ) 

61 coadds = cT.Input( 

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

63 name="{inputCoaddName}Coadd_calexp", 

64 storageClass="ExposureF", 

65 multiple=True, 

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

67 ) 

68 coadds_cell = cT.Input( 

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

70 name="{inputCoaddName}CoaddCell", 

71 storageClass="MultipleCellCoadd", 

72 multiple=True, 

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

74 ) 

75 backgrounds = cT.Input( 

76 doc="Background model to subtract from the cell-based coadd", 

77 name="{inputCoaddName}Coadd_calexp_background", 

78 storageClass="Background", 

79 multiple=True, 

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

81 ) 

82 deconvolvedCoadds = cT.Input( 

83 doc="Deconvolved coadds", 

84 name="deconvolved_{inputCoaddName}_coadd", 

85 storageClass="ExposureF", 

86 multiple=True, 

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

88 ) 

89 outputSchema = cT.InitOutput( 

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

91 name="{outputCoaddName}Coadd_deblendedFlux_schema", 

92 storageClass="SourceCatalog" 

93 ) 

94 deblendedCatalog = cT.Output( 

95 doc="Catalogs produced by multiband deblending", 

96 name="{outputCoaddName}Coadd_deblendedCatalog", 

97 storageClass="SourceCatalog", 

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

99 ) 

100 scarletModelData = cT.Output( 

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

102 name="{outputCoaddName}Coadd_scarletModelData", 

103 storageClass="LsstScarletModelData", 

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

105 ) 

106 objectParents = cT.Output( 

107 doc="Parents of the deblended objects", 

108 name="object_parent_patch", 

109 storageClass="SourceCatalog", 

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

111 ) 

112 

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

114 super().__init__(config=config) 

115 if config: 

116 if config.useCellCoadds: 

117 del self.coadds 

118 else: 

119 del self.coadds_cell 

120 del self.backgrounds 

121 

122 

123class DeblendCoaddSourcesMultiConfig(PipelineTaskConfig, 

124 pipelineConnections=DeblendCoaddSourcesMultiConnections): 

125 useCellCoadds = Field[bool]( 

126 doc="Use cell-based coadds instead of regular coadds?", 

127 default=False, 

128 ) 

129 multibandDeblend = ConfigurableField( 

130 target=ScarletDeblendTask, 

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

132 ) 

133 idGenerator = SkyMapIdGeneratorConfig.make_field() 

134 

135 

136class DeblendCoaddSourcesMultiTask(PipelineTask): 

137 ConfigClass = DeblendCoaddSourcesMultiConfig 

138 _DefaultName = "deblendCoaddSourcesMulti" 

139 

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

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

142 schema = initInputs["inputSchema"].schema 

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

144 self.schemaMapper = afwTable.SchemaMapper(schema) 

145 self.schemaMapper.addMinimalSchema(schema) 

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

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

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

149 

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

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

152 # all input lists to match this band order. 

153 # Note: sometimes deconvolution fails. If this happens then 

154 # the dataIds missing from deconvolvedRefs will be removed 

155 # during the process. 

156 deconvolvedRefs = inputRefs.deconvolvedCoadds 

157 bandOrder = [dRef.dataId["band"] for dRef in deconvolvedRefs] 

158 bandOrder.sort() 

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

160 inputs = butlerQC.get(inputRefs) 

161 bands = [dRef.dataId["band"] for dRef in deconvolvedRefs] 

162 mergedDetections = inputs.pop("mergedDetections") 

163 if self.config.useCellCoadds: 

164 exposures = [mcc.stitch().asExposure() for mcc in inputs.pop("coadds_cell")] 

165 backgrounds = inputs.pop("backgrounds") 

166 for exposure, background in zip(exposures, backgrounds): 

167 exposure.image -= background.getImage() 

168 coadds = exposures 

169 else: 

170 coadds = inputs.pop("coadds") 

171 

172 # Ensure that the coadd bands and deconvolved coadd bands match 

173 coaddRefs = inputRefs.coadds_cell if self.config.useCellCoadds else inputRefs.coadds 

174 coaddBands = [dRef.dataId["band"] for dRef in coaddRefs] 

175 if bands != coaddBands: 

176 self.log.error("Coadd bands %s != deconvolved coadd bands %s", bands, coaddBands) 

177 raise RuntimeError( 

178 "Number of coadd bands and deconvolved coadd bands do not match. " 

179 "This should never happen and indicates a bug in reorderRefs." 

180 ) 

181 

182 deconvolvedCoadds = inputs.pop("deconvolvedCoadds") 

183 

184 # Check that all inputs have been extracted correctly. 

185 assert not inputs, "runQuantum got extra inputs" 

186 

187 outputs = self.run( 

188 coadds=coadds, 

189 bands=bands, 

190 mergedDetections=mergedDetections, 

191 idFactory=self.config.idGenerator.apply(butlerQC.quantum.dataId).make_table_id_factory(), 

192 deconvolvedCoadds=deconvolvedCoadds, 

193 ) 

194 butlerQC.put(outputs, outputRefs) 

195 

196 def run(self, coadds, bands, mergedDetections, deconvolvedCoadds, idFactory): 

197 sources = self._makeSourceCatalog(mergedDetections, idFactory) 

198 multiExposure = afwImage.MultibandExposure.fromExposures(bands, coadds) 

199 mDeconvolved = afwImage.MultibandExposure.fromExposures(bands, deconvolvedCoadds) 

200 result = self.multibandDeblend.run(multiExposure, mDeconvolved, sources) 

201 return result 

202 

203 def _makeSourceCatalog(self, mergedDetections, idFactory): 

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

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

206 # with the largest id already in the catalog. 

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

208 idFactory.notify(maxId) 

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

210 sources = afwTable.SourceCatalog(table) 

211 sources.extend(mergedDetections, self.schemaMapper) 

212 return sources