Coverage for python/lsst/pipe/tasks/match_tract_catalog.py: 62%

58 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-28 03:53 -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__ = [ 

23 'MatchTractCatalogSubConfig', 'MatchTractCatalogSubTask', 

24 'MatchTractCatalogConfig', 'MatchTractCatalogTask' 

25] 

26 

27import lsst.afw.geom as afwGeom 

28import lsst.pex.config as pexConfig 

29import lsst.pipe.base as pipeBase 

30import lsst.pipe.base.connectionTypes as cT 

31from lsst.skymap import BaseSkyMap 

32 

33from abc import ABC, abstractmethod 

34 

35import pandas as pd 

36from typing import Tuple, Set 

37 

38 

39MatchTractCatalogBaseTemplates = { 

40 "name_input_cat_ref": "truth_summary", 

41 "name_input_cat_target": "objectTable_tract", 

42} 

43 

44 

45class MatchTractCatalogConnections( 

46 pipeBase.PipelineTaskConnections, 

47 dimensions=("tract", "skymap"), 

48 defaultTemplates=MatchTractCatalogBaseTemplates, 

49): 

50 cat_ref = cT.Input( 

51 doc="Reference object catalog to match from", 

52 name="{name_input_cat_ref}", 

53 storageClass="DataFrame", 

54 dimensions=("tract", "skymap"), 

55 deferLoad=True, 

56 ) 

57 cat_target = cT.Input( 

58 doc="Target object catalog to match", 

59 name="{name_input_cat_target}", 

60 storageClass="DataFrame", 

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

62 deferLoad=True, 

63 ) 

64 skymap = cT.Input( 

65 doc="Input definition of geometry/bbox and projection/wcs for coadded exposures", 

66 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

67 storageClass="SkyMap", 

68 dimensions=("skymap",), 

69 ) 

70 cat_output_ref = cT.Output( 

71 doc="Reference matched catalog with indices of target matches", 

72 name="match_ref_{name_input_cat_ref}_{name_input_cat_target}", 

73 storageClass="DataFrame", 

74 dimensions=("tract", "skymap"), 

75 ) 

76 cat_output_target = cT.Output( 

77 doc="Target matched catalog with indices of reference matches", 

78 name="match_target_{name_input_cat_ref}_{name_input_cat_target}", 

79 storageClass="DataFrame", 

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

81 ) 

82 

83 

84class MatchTractCatalogSubConfig(pexConfig.Config): 

85 """Config class for the MatchTractCatalogSubTask to define methods returning 

86 values that depend on multiple config settings. 

87 """ 

88 @property 

89 @abstractmethod 

90 def columns_in_ref(self) -> Set[str]: 

91 raise NotImplementedError() 

92 

93 @property 

94 @abstractmethod 

95 def columns_in_target(self) -> Set[str]: 

96 raise NotImplementedError() 

97 

98 

99class MatchTractCatalogSubTask(pipeBase.Task, ABC): 

100 """An abstract interface for subtasks of MatchTractCatalogTask to match 

101 two tract object catalogs. 

102 

103 Parameters 

104 ---------- 

105 **kwargs 

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

107 constructor. 

108 """ 

109 ConfigClass = MatchTractCatalogSubConfig 

110 

111 def __init__(self, **kwargs): 

112 super().__init__(**kwargs) 

113 

114 @abstractmethod 

115 def run( 

116 self, 

117 catalog_ref: pd.DataFrame, 

118 catalog_target: pd.DataFrame, 

119 wcs: afwGeom.SkyWcs = None, 

120 ) -> pipeBase.Struct: 

121 """Match sources in a reference tract catalog with a target catalog. 

122 

123 Parameters 

124 ---------- 

125 catalog_ref : `pandas.DataFrame` 

126 A reference catalog to match objects/sources from. 

127 catalog_target : `pandas.DataFrame` 

128 A target catalog to match reference objects/sources to. 

129 wcs : `lsst.afw.image.SkyWcs` 

130 A coordinate system to convert catalog positions to sky coordinates. 

131 

132 Returns 

133 ------- 

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

135 A struct with output_ref and output_target attribute containing the 

136 output matched catalogs. 

137 """ 

138 raise NotImplementedError() 

139 

140 

141class MatchTractCatalogConfig( 

142 pipeBase.PipelineTaskConfig, 

143 pipelineConnections=MatchTractCatalogConnections, 

144): 

145 """Configure a MatchTractCatalogTask, including a configurable matching subtask. 

146 """ 

147 match_tract_catalog = pexConfig.ConfigurableField( 

148 target=MatchTractCatalogSubTask, 

149 doc="Task to match sources in a reference tract catalog with a target catalog", 

150 ) 

151 

152 def get_columns_in(self) -> Tuple[Set, Set]: 

153 """Get the set of input columns required for matching. 

154 

155 Returns 

156 ------- 

157 columns_ref : `set` [`str`] 

158 The set of required input catalog column names. 

159 columns_target : `set` [`str`] 

160 The set of required target catalog column names. 

161 """ 

162 try: 

163 columns_ref, columns_target = (self.match_tract_catalog.columns_in_ref, 

164 self.match_tract_catalog.columns_in_target) 

165 except AttributeError as err: 

166 raise RuntimeError(f'{__class__}.match_tract_catalog must have columns_in_ref and' 

167 f' columns_in_target attributes: {err}') from None 

168 return set(columns_ref), set(columns_target) 

169 

170 

171class MatchTractCatalogTask(pipeBase.PipelineTask): 

172 """Match sources in a reference tract catalog with those in a target catalog. 

173 """ 

174 ConfigClass = MatchTractCatalogConfig 

175 _DefaultName = "MatchTractCatalog" 

176 

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

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

179 self.makeSubtask("match_tract_catalog") 

180 

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

182 inputs = butlerQC.get(inputRefs) 

183 columns_ref, columns_target = self.config.get_columns_in() 

184 skymap = inputs.pop("skymap") 

185 

186 outputs = self.run( 

187 catalog_ref=inputs['cat_ref'].get(parameters={'columns': columns_ref}), 

188 catalog_target=inputs['cat_target'].get(parameters={'columns': columns_target}), 

189 wcs=skymap[butlerQC.quantum.dataId["tract"]].wcs, 

190 ) 

191 butlerQC.put(outputs, outputRefs) 

192 

193 def run( 

194 self, 

195 catalog_ref: pd.DataFrame, 

196 catalog_target: pd.DataFrame, 

197 wcs: afwGeom.SkyWcs = None, 

198 ) -> pipeBase.Struct: 

199 """Match sources in a reference tract catalog with a target catalog. 

200 

201 Parameters 

202 ---------- 

203 catalog_ref : `pandas.DataFrame` 

204 A reference catalog to match objects/sources from. 

205 catalog_target : `pandas.DataFrame` 

206 A target catalog to match reference objects/sources to. 

207 wcs : `lsst.afw.image.SkyWcs` 

208 A coordinate system to convert catalog positions to sky coordinates, 

209 if necessary. 

210 

211 Returns 

212 ------- 

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

214 A struct with output_ref and output_target attribute containing the 

215 output matched catalogs. 

216 """ 

217 output = self.match_tract_catalog.run(catalog_ref, catalog_target, wcs=wcs) 

218 if output.exceptions: 

219 self.log.warn('Exceptions: %s', output.exceptions) 

220 retStruct = pipeBase.Struct(cat_output_ref=output.cat_output_ref, 

221 cat_output_target=output.cat_output_target) 

222 return retStruct