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 

33 

34__all__ = ["MatchApFakesTask", 

35 "MatchApFakesConfig", 

36 "MatchApFakesConnections"] 

37 

38 

39class MatchApFakesConnections(PipelineTaskConnections, 

40 defaultTemplates={"coaddName": "deep", 

41 "fakesType": "fakes_"}, 

42 dimensions=("tract", 

43 "skymap", 

44 "instrument", 

45 "visit", 

46 "detector")): 

47 fakeCat = connTypes.Input( 

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

49 name="{fakesType}fakeSourceCat", 

50 storageClass="DataFrame", 

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

52 ) 

53 diffIm = connTypes.Input( 

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

55 name="{fakesType}{coaddName}Diff_differenceExp", 

56 storageClass="ExposureF", 

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

58 ) 

59 associatedDiaSources = connTypes.Input( 

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

61 "SDMification.", 

62 name="{fakesType}{coaddName}Diff_assocDiaSrc", 

63 storageClass="DataFrame", 

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

65 ) 

66 matchedDiaSources = connTypes.Output( 

67 doc="", 

68 name="{fakesType}{coaddName}Diff_matchDiaSrc", 

69 storageClass="DataFrame", 

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

71 ) 

72 

73 

74class MatchApFakesConfig( 

75 InsertFakesConfig, 

76 pipelineConnections=MatchApFakesConnections): 

77 """Config for MatchApFakesTask. 

78 """ 

79 matchDistanceArcseconds = pexConfig.RangeField( 

80 doc="Distance in arcseconds to ", 

81 dtype=float, 

82 default=0.5, 

83 min=0, 

84 max=10, 

85 ) 

86 

87 

88class MatchApFakesTask(PipelineTask): 

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

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

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

92 template exposure. 

93 """ 

94 

95 _DefaultName = "matchApFakes" 

96 ConfigClass = MatchApFakesConfig 

97 

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

99 inputs = butlerQC.get(inputRefs) 

100 

101 outputs = self.run(**inputs) 

102 butlerQC.put(outputs, outputRefs) 

103 

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

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

106 

107 Parameters 

108 ---------- 

109 fakeCat : `pandas.DataFrame` 

110 Catalog of fakes to match to detected diaSources. 

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

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

113 associatedDiaSources : `pandas.DataFrame` 

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

115 

116 Returns 

117 ------- 

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

119 Results struct with components. 

120 

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

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

123 """ 

124 trimmedFakes = self._trimFakeCat(fakeCat, diffIm) 

125 nPossibleFakes = len(trimmedFakes) 

126 

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

128 trimmedFakes[self.config.decColName]) 

129 diaSrcVects = self._getVectors( 

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

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

132 

133 diaSrcTree = cKDTree(diaSrcVects) 

134 dist, idxs = diaSrcTree.query( 

135 fakeVects, 

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

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

138 

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

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

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

142 

143 return Struct( 

144 matchedDiaSources=matchedFakes.merge( 

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

146 ) 

147 

148 def _trimFakeCat(self, fakeCat, image): 

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

150 

151 Parameters 

152 ---------- 

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

154 The catalog of fake sources to be input 

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

156 The image into which the fake sources should be added 

157 

158 Returns 

159 ------- 

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

161 The original fakeCat trimmed to the area of the image 

162 """ 

163 wcs = image.getWcs() 

164 

165 bbox = Box2D(image.getBBox()) 

166 

167 def trim(row): 

168 coord = SpherePoint(row[self.config.raColName], 

169 row[self.config.decColName], 

170 radians) 

171 cent = wcs.skyToPixel(coord) 

172 return bbox.contains(cent) 

173 

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

175 

176 def _getVectors(self, ras, decs): 

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

178 

179 Parameters 

180 ---------- 

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

182 RA coordinates in radians. 

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

184 Dec coordinates in radians. 

185 

186 Returns 

187 ------- 

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

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

190 """ 

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

192 

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

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

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

196 

197 return vectors