1# This file is part of ap_pipe. 

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"""Methods to match an input catalog to a set of fakes in AP. 

23""" 

24 

25import numpy as np 

26from scipy.spatial import cKDTree 

27 

28from lsst.geom import Box2D, radians, SpherePoint 

29import lsst.pex.config as pexConfig 

30from lsst.pipe.base import PipelineTask, PipelineTaskConnections, Struct 

31import lsst.pipe.base.connectionTypes as connTypes 

32from lsst.pipe.tasks.insertFakes import InsertFakesConfig 

33from lsst.sphgeom import ConvexPolygon 

34 

35__all__ = ["MatchApFakesTask", 

36 "MatchApFakesConfig", 

37 "MatchApFakesConnections"] 

38 

39 

40class MatchApFakesConnections(PipelineTaskConnections, 

41 defaultTemplates={"CoaddName": "deep", 

42 "fakesType": ""}, 

43 dimensions=("tract", 

44 "skymap", 

45 "instrument", 

46 "visit", 

47 "detector")): 

48 fakeCat = connTypes.Input( 

49 doc="Catalog of fake sources to draw inputs from.", 

50 name="{CoaddName}Coadd_fakeSourceCat", 

51 storageClass="DataFrame", 

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

53 ) 

54 diffIm = connTypes.Input( 

55 doc="Difference image on which the DiaSources were detected.", 

56 name="{fakesType}{CoaddName}Diff_differenceExp", 

57 storageClass="ExposureF", 

58 dimensions=("instrument", "visit", "detector"), 

59 ) 

60 associatedDiaSources = connTypes.Input( 

61 doc="Optional output storing the DiaSource catalog after matching and " 

62 "SDMification.", 

63 name="{fakesType}{CoaddName}Diff_assocDiaSrc", 

64 storageClass="DataFrame", 

65 dimensions=("instrument", "visit", "detector"), 

66 ) 

67 matchedDiaSources = connTypes.Output( 

68 doc="", 

69 name="{fakesType}{CoaddName}Diff_matchDiaSrc", 

70 storageClass="DataFrame", 

71 dimensions=("instrument", "visit", "detector"), 

72 ) 

73 

74 

75class MatchApFakesConfig( 

76 InsertFakesConfig, 

77 pipelineConnections=MatchApFakesConnections): 

78 """Config for MatchApFakesTask. 

79 """ 

80 matchDistanceArcseconds = pexConfig.RangeField( 

81 doc="Distance in arcseconds to ", 

82 dtype=float, 

83 default=1, 

84 min=0, 

85 max=10, 

86 ) 

87 

88 

89class MatchApFakesTask(PipelineTask): 

90 """Create and store a set of spatially uniform star fakes over the sphere 

91 for use in AP processing. Additionally assign random magnitudes to said 

92 fakes and assign them to be inserted into either a visit exposure or 

93 template exposure. 

94 """ 

95 

96 _DefaultName = "matchApFakes" 

97 ConfigClass = MatchApFakesConfig 

98 

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

100 inputs = butlerQC.get(inputRefs) 

101 

102 outputs = self.run(**inputs) 

103 butlerQC.put(outputs, outputRefs) 

104 

105 def run(self, fakeCat, diffIm, associatedDiaSources): 

106 """Match fakes to detected diaSources within a difference image bound. 

107 

108 Parameters 

109 ---------- 

110 fakeCat : `pandas.DataFrame` 

111 Catalog of fakes to match to detected diaSources. 

112 diffIm : `lsst.afw.image.Exposure` 

113 Difference image where ``associatedDiaSources`` were detected in. 

114 associatedDiaSources : `pandas.DataFrame` 

115 Catalog of difference image sources detected in ``diffIm``. 

116 

117 Returns 

118 ------- 

119 result : `lsst.pipe.base.Struct` 

120 Results struct with components. 

121 

122 - ``matchedDiaSources`` : Fakes matched to input diaSources. Has 

123 length of ``fakeCat``. (`pandas.DataFrame`) 

124 """ 

125 trimmedFakes = self._trimFakeCat(fakeCat, diffIm) 

126 nPossibleFakes = len(trimmedFakes) 

127 

128 fakeVects = self._getVectors(trimmedFakes[self.config.raColName], 

129 trimmedFakes[self.config.decColName]) 

130 diaSrcVects = self._getVectors( 

131 np.radians(associatedDiaSources.loc[:, "ra"]), 

132 np.radians(associatedDiaSources.loc[:, "decl"])) 

133 

134 diaSrcTree = cKDTree(diaSrcVects) 

135 dist, idxs = diaSrcTree.query( 

136 fakeVects, 

137 distance_upper_bound=np.radians(self.config.matchDistanceArcseconds / 3600)) 

138 nFakesFound = np.isfinite(dist).sum() 

139 

140 self.log.info(f"Found {nFakesFound} out of {nPossibleFakes} possible.") 

141 diaSrcIds = associatedDiaSources.iloc[np.where(np.isfinite(dist), idxs, 0)]["diaSourceId"].to_numpy() 

142 matchedFakes = trimmedFakes.assign(diaSourceId=np.where(np.isfinite(dist), diaSrcIds, 0)) 

143 

144 return Struct( 

145 matchedDiaSources=matchedFakes.merge( 

146 associatedDiaSources.reset_index(drop=True), on="diaSourceId", how="left") 

147 ) 

148 

149 def _trimFakeCat(self, fakeCat, image): 

150 """Trim the fake cat to about the size of the input image. 

151 

152 Parameters 

153 ---------- 

154 fakeCat : `pandas.core.frame.DataFrame` 

155 The catalog of fake sources to be input 

156 image : `lsst.afw.image.exposure.exposure.ExposureF` 

157 The image into which the fake sources should be added 

158 

159 Returns 

160 ------- 

161 fakeCat : `pandas.core.frame.DataFrame` 

162 The original fakeCat trimmed to the area of the image 

163 """ 

164 wcs = image.getWcs() 

165 

166 bbox = Box2D(image.getBBox()) 

167 corners = bbox.getCorners() 

168 

169 skyCorners = wcs.pixelToSky(corners) 

170 region = ConvexPolygon([s.getVector() for s in skyCorners]) 

171 

172 def trim(row): 

173 coord = SpherePoint(row[self.config.raColName], row[self.config.decColName], radians) 

174 return region.contains(coord.getVector()) 

175 

176 return fakeCat[fakeCat.apply(trim, axis=1)] 

177 

178 def _getVectors(self, ras, decs): 

179 """Convert ra dec to unit vectors on the sphere. 

180 

181 Parameters 

182 ---------- 

183 ras : `numpy.ndarray`, (N,) 

184 RA coordinates in radians. 

185 decs : `numpy.ndarray`, (N,) 

186 Dec coordinates in radians. 

187 

188 Returns 

189 ------- 

190 vectors : `numpy.ndarray`, (N, 3) 

191 Vectors on the unit sphere for the given RA/DEC values. 

192 """ 

193 vectors = np.empty((len(ras), 3)) 

194 

195 vectors[:, 2] = np.sin(decs) 

196 vectors[:, 0] = np.cos(decs) * np.cos(ras) 

197 vectors[:, 1] = np.cos(decs) * np.sin(ras) 

198 

199 return vectors