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

53 statements  

« prev     ^ index     » next       coverage.py v6.4.2, created at 2022-07-19 05:42 -0700

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 

29from astropy import units as u 

30 

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 else: 

109 self.log.debug("Matching solar system objects:\n%s", maskedObjects) 

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

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

112 

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

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

115 diaSourceCatalog["decl"]) 

116 

117 # Create KDTree of DIA sources 

118 tree = cKDTree(vectors) 

119 

120 nFound = 0 

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

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

123 # fancier later. 

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

125 

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

127 

128 # Which DIA Sources fall within r? 

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

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

131 nFound += 1 

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

133 

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

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

136 return pipeBase.Struct( 

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

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

139 nTotalSsObjects=nSolarSystemObjects, 

140 nAssociatedSsObjects=nFound) 

141 

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

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

144 bounding box. 

145 

146 Parameters 

147 ---------- 

148 solarSystemObjects : `pandas.DataFrame` 

149 SolarSystemObjects to mask to ``exposure``. 

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

151 Exposure to mask to. 

152 marginArcsec : `float` 

153 Maximum possible matching radius to pad onto the exposure bounding 

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

155 be used. 

156 

157 Returns 

158 ------- 

159 maskedSolarSystemObjects : `pandas.DataFrame` 

160 Set of SolarSystemObjects contained within the exposure bounds. 

161 """ 

162 wcs = exposure.getWcs() 

163 padding = min( 

164 int(np.ceil(marginArcsec / wcs.getPixelScale().asArcseconds())), 

165 self.config.maxPixelMargin) 

166 

167 return solarSystemObjects[exposure.containsSkyCoords( 

168 solarSystemObjects['ra'].to_numpy() * u.degree, 

169 solarSystemObjects['decl'].to_numpy() * u.degree, 

170 padding)] 

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