Coverage for python/lsst/ap/association/ssoAssociation.py: 35%

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

56 statements  

1# This file is part of ap_association. 

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"""Spatial association for Solar System Objects.""" 

23 

24__all__ = ["SolarSystemAssociationConfig", "SolarSystemAssociationTask"] 

25 

26import numpy as np 

27import pandas as pd 

28from scipy.spatial import cKDTree 

29 

30import lsst.geom as geom 

31import lsst.pex.config as pexConfig 

32import lsst.pipe.base as pipeBase 

33from lsst.utils.timer import timeMethod 

34 

35 

36class SolarSystemAssociationConfig(pexConfig.Config): 

37 """Config class for SolarSystemAssociationTask. 

38 """ 

39 maxDistArcSeconds = pexConfig.Field( 

40 dtype=float, 

41 doc='Maximum distance in arcseconds to test for a DIASource to be a ' 

42 'match to a SSObject.', 

43 default=2.0, 

44 ) 

45 maxPixelMargin = pexConfig.RangeField( 

46 doc="Maximum padding to add to the ccd bounding box before masking " 

47 "SolarSystem objects to the ccd footprint. The bounding box will " 

48 "be padded by the minimum of this number or the max uncertainty " 

49 "of the SolarSystemObjects in pixels.", 

50 dtype=int, 

51 default=100, 

52 min=0, 

53 ) 

54 

55 

56class SolarSystemAssociationTask(pipeBase.Task): 

57 """Associate DIASources into existing SolarSystem Objects. 

58 

59 This task performs the association of detected DIASources in a visit 

60 with known solar system objects. 

61 """ 

62 ConfigClass = SolarSystemAssociationConfig 

63 _DefaultName = "ssoAssociation" 

64 

65 @timeMethod 

66 def run(self, diaSourceCatalog, solarSystemObjects, exposure): 

67 """Create a searchable tree of unassociated DiaSources and match 

68 to the nearest ssoObject. 

69 

70 Parameters 

71 ---------- 

72 diaSourceCatalog : `pandas.DataFrame` 

73 Catalog of DiaSources. Modified in place to add ssObjectId to 

74 successfully associated DiaSources. 

75 solarSystemObjects : `pandas.DataFrame` 

76 Set of solar system objects that should be within the footprint 

77 of the current visit. 

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

79 Exposure where the DiaSources in ``diaSourceCatalog`` were 

80 detected in. 

81 

82 Returns 

83 ------- 

84 resultsStruct : `lsst.pipe.base.Struct` 

85 

86 - ``ssoAssocDiaSources`` : DiaSources that were associated with 

87 solar system objects in this visit. (`pandas.DataFrame`) 

88 - ``unAssocDiaSources`` : Set of DiaSources that were not 

89 associated with any solar system object. (`pandas.DataFrame`) 

90 - ``nTotalSsObjects`` : Total number of SolarSystemObjects 

91 contained in the CCD footprint. (`int`) 

92 - ``nAssociatedSsObjects`` : Number of SolarSystemObjects 

93 that were associated with DiaSources. 

94 """ 

95 maskedObjects = self._maskToCcdRegion( 

96 solarSystemObjects, 

97 exposure, 

98 solarSystemObjects["Err(arcsec)"].max()) 

99 nSolarSystemObjects = len(maskedObjects) 

100 if nSolarSystemObjects <= 0: 

101 self.log.info("No SolarSystemObjects found in detector bounding " 

102 "box.") 

103 return pipeBase.Struct( 

104 ssoAssocDiaSources=pd.DataFrame(columns=diaSourceCatalog.columns), 

105 unAssocDiaSources=diaSourceCatalog, 

106 nTotalSsObjects=0, 

107 nAssociatedSsObjects=0) 

108 self.log.info("Attempting to associate %d...", nSolarSystemObjects) 

109 maxRadius = np.deg2rad(self.config.maxDistArcSeconds / 3600) 

110 

111 # Transform DIA RADEC coordinates to unit sphere xyz for tree building. 

112 vectors = self._radec_to_xyz(diaSourceCatalog["ra"], 

113 diaSourceCatalog["decl"]) 

114 

115 # Create KDTree of DIA sources 

116 tree = cKDTree(vectors) 

117 

118 nFound = 0 

119 # Query the KDtree for DIA nearest neighbors to SSOs. Currently only 

120 # picks the DiaSource with the shortest distance. We can do something 

121 # fancier later. 

122 for index, ssObject in maskedObjects.iterrows(): 

123 

124 ssoVect = self._radec_to_xyz(ssObject["ra"], ssObject["decl"]) 

125 

126 # Which DIA Sources fall within r? 

127 dist, idx = tree.query(ssoVect, distance_upper_bound=maxRadius) 

128 if np.isfinite(dist[0]): 

129 nFound += 1 

130 diaSourceCatalog.loc[idx[0], "ssObjectId"] = ssObject["ssObjectId"] 

131 

132 self.log.info("Successfully associated %d SolarSystemObjects.", nFound) 

133 assocMask = diaSourceCatalog["ssObjectId"] != 0 

134 return pipeBase.Struct( 

135 ssoAssocDiaSources=diaSourceCatalog[assocMask].reset_index(drop=True), 

136 unAssocDiaSources=diaSourceCatalog[~assocMask].reset_index(drop=True), 

137 nTotalSsObjects=nSolarSystemObjects, 

138 nAssociatedSsObjects=nFound) 

139 

140 def _maskToCcdRegion(self, solarSystemObjects, exposure, marginArcsec): 

141 """Mask the input SolarSystemObjects to only those in the exposure 

142 bounding box. 

143 

144 Parameters 

145 ---------- 

146 solarSystemObjects : `pandas.DataFrame` 

147 SolarSystemObjects to mask to ``exposure``. 

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

149 Exposure to mask to. 

150 marginArcsec : `float` 

151 Maximum possible matching radius to pad onto the exposure bounding 

152 box. If greater than ``maxPixelMargin``, ``maxPixelMargin`` will 

153 be used. 

154 

155 Returns 

156 ------- 

157 maskedSolarSystemObjects : `pandas.DataFrame` 

158 Set of SolarSystemObjects contained within the exposure bounds. 

159 """ 

160 wcs = exposure.getWcs() 

161 padding = min( 

162 int(np.ceil(wcs.getPixelScale().asArcseconds()*marginArcsec)), 

163 self.config.maxPixelMargin) 

164 bbox = geom.Box2D(exposure.getBBox()) 

165 bbox.grow(padding) 

166 

167 mapping = wcs.getTransform().getMapping() 

168 x, y = mapping.applyInverse( 

169 np.radians(solarSystemObjects[['ra', 'decl']].T.to_numpy())) 

170 return solarSystemObjects[bbox.contains(x, y)] 

171 

172 def _radec_to_xyz(self, ras, decs): 

173 """Convert input ra/dec coordinates to spherical unit-vectors. 

174 

175 Parameters 

176 ---------- 

177 ras : `array-like` 

178 RA coordinates of objects in degrees. 

179 decs : `array-like` 

180 DEC coordinates of objects in degrees. 

181 

182 Returns 

183 ------- 

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

185 Output unit-vectors 

186 """ 

187 ras = np.radians(ras) 

188 decs = np.radians(decs) 

189 try: 

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

191 except TypeError: 

192 vectors = np.empty((1, 3)) 

193 

194 sin_dec = np.sin(np.pi / 2 - decs) 

195 vectors[:, 0] = sin_dec * np.cos(ras) 

196 vectors[:, 1] = sin_dec * np.sin(ras) 

197 vectors[:, 2] = np.cos(np.pi / 2 - decs) 

198 

199 return vectors