Coverage for python/lsst/pipe/tasks/ssoAssociation.py: 17%

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 

26from astropy import units as u 

27from astropy.table import Table, Column 

28from astropy.coordinates import SkyCoord 

29import healpy as hp 

30import numpy as np 

31from numpy.polynomial.chebyshev import Chebyshev, chebval 

32from scipy.spatial import cKDTree 

33 

34from lsst.afw.image.exposure.exposureUtils import bbox_contains_sky_coords 

35import lsst.pex.config as pexConfig 

36import lsst.pipe.base as pipeBase 

37from lsst.utils.timer import timeMethod 

38 

39 

40class SolarSystemAssociationConfig(pexConfig.Config): 

41 """Config class for SolarSystemAssociationTask. 

42 """ 

43 maxDistArcSeconds = pexConfig.Field( 

44 dtype=float, 

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

46 'match to a SSObject.', 

47 default=1.0, 

48 ) 

49 maxPixelMargin = pexConfig.RangeField( 

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

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

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

53 "of the SolarSystemObjects in pixels.", 

54 dtype=int, 

55 default=100, 

56 min=0, 

57 ) 

58 

59 

60class SolarSystemAssociationTask(pipeBase.Task): 

61 """Associate DIASources into existing SolarSystem Objects. 

62 

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

64 with known solar system objects. 

65 """ 

66 ConfigClass = SolarSystemAssociationConfig 

67 _DefaultName = "ssoAssociation" 

68 

69 @timeMethod 

70 def run(self, diaSourceCatalog, solarSystemObjects, visitInfo, bbox, wcs): 

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

72 to the nearest ssoObject. 

73 

74 Parameters 

75 ---------- 

76 diaSourceCatalog : `astropy.table.Table` 

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

78 successfully associated DiaSources. 

79 solarSystemObjects : `astropy.table.Table` 

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

81 of the current visit. 

82 visitInfo : `lsst.afw.image.VisitInfo` 

83 visitInfo of exposure used for exposure time 

84 bbox : `lsst.geom.Box2I` 

85 bbox of exposure used for masking 

86 wcs : `lsst.afw.geom.SkyWcs` 

87 wcs of exposure used for masking 

88 

89 Returns 

90 ------- 

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

92 

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

94 solar system objects in this visit. (`astropy.table.Table`) 

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

96 associated with any solar system object. (`astropy.table.Table`) 

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

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

99 - ``nAssociatedSsObjects`` : Number of SolarSystemObjects 

100 that were associated with DiaSources. (`int`) 

101 - ``ssSourceData`` : ssSource table data. (`Astropy.table.Table`) 

102 """ 

103 

104 nSolarSystemObjects = len(solarSystemObjects) 

105 if nSolarSystemObjects <= 0: 

106 return self._return_empty(diaSourceCatalog, solarSystemObjects) 

107 

108 mjd_midpoint = visitInfo.date.toAstropy().tai.mjd 

109 ref_time = mjd_midpoint - solarSystemObjects["tmin"].value[0] # all tmin should be identical 

110 

111 solarSystemObjects['obs_position'] = [ 

112 np.array([chebval(ref_time, row['obs_x_poly']), 

113 chebval(ref_time, row['obs_y_poly']), 

114 chebval(ref_time, row['obs_z_poly'])]) 

115 for row in solarSystemObjects] 

116 solarSystemObjects['obs_velocity'] = [ 

117 np.array([chebval(ref_time, Chebyshev(row['obs_x_poly']).deriv().coef), 

118 chebval(ref_time, Chebyshev(row['obs_y_poly']).deriv().coef), 

119 chebval(ref_time, Chebyshev(row['obs_z_poly']).deriv().coef)]) 

120 for row in solarSystemObjects] 

121 solarSystemObjects['obj_position'] = [ 

122 np.array([chebval(ref_time, row['obj_x_poly']), 

123 chebval(ref_time, row['obj_y_poly']), 

124 chebval(ref_time, row['obj_z_poly'])]) 

125 for row in solarSystemObjects] 

126 solarSystemObjects['obj_velocity'] = [ 

127 np.array([chebval(ref_time, Chebyshev(row['obj_x_poly']).deriv().coef), 

128 chebval(ref_time, Chebyshev(row['obj_y_poly']).deriv().coef), 

129 chebval(ref_time, Chebyshev(row['obj_z_poly']).deriv().coef)]) 

130 for row in solarSystemObjects] 

131 vector = np.vstack(solarSystemObjects['obj_position'].value 

132 - solarSystemObjects['obs_position'].value) 

133 ras, decs = np.vstack(hp.vec2ang(vector, lonlat=True)) 

134 solarSystemObjects['ra'] = ras 

135 solarSystemObjects['dec'] = decs 

136 solarSystemObjects['obs_position_x'], solarSystemObjects['obs_position_y'], \ 

137 solarSystemObjects['obs_position_z'] = solarSystemObjects['obs_position'].value.T 

138 solarSystemObjects['heliocentricX'], solarSystemObjects['heliocentricY'], \ 

139 solarSystemObjects['heliocentricZ'] = solarSystemObjects['obj_position'].value.T 

140 solarSystemObjects['obs_velocity_x'], solarSystemObjects['obs_velocity_y'], \ 

141 solarSystemObjects['obs_velocity_z'] = solarSystemObjects['obs_velocity'].value.T 

142 solarSystemObjects['heliocentricVX'], solarSystemObjects['heliocentricVY'], \ 

143 solarSystemObjects['heliocentricVZ'] = solarSystemObjects['obj_velocity'].value.T 

144 solarSystemObjects['topocentric_position'], solarSystemObjects['topocentric_velocity'] = ( 

145 solarSystemObjects['obj_position'] - solarSystemObjects['obs_position'], 

146 solarSystemObjects['obj_velocity'] - solarSystemObjects['obs_velocity'], 

147 ) 

148 solarSystemObjects['topocentricX'], solarSystemObjects['topocentricY'], \ 

149 solarSystemObjects['topocentricZ'] = ( 

150 np.array(list(solarSystemObjects['topocentric_position'].value)).T 

151 ) 

152 solarSystemObjects['topocentricVX'], solarSystemObjects['topocentricVY'], \ 

153 solarSystemObjects['topocentricVZ'] = ( 

154 np.array(list(solarSystemObjects['topocentric_velocity'].value)).T 

155 ) 

156 solarSystemObjects['heliocentricVX'], solarSystemObjects['heliocentricVY'], \ 

157 solarSystemObjects['heliocentricVZ'] = np.array(list(solarSystemObjects['obj_velocity'].value)).T 

158 solarSystemObjects['heliocentricDist'] = np.linalg.norm(solarSystemObjects['obj_position'], axis=1) 

159 solarSystemObjects['topocentricDist'] = np.linalg.norm(solarSystemObjects['topocentric_position'], 

160 axis=1) 

161 solarSystemObjects['phaseAngle'] = np.degrees(np.arccos(np.sum( 

162 solarSystemObjects['obj_position'].T * solarSystemObjects['topocentric_position'].T 

163 / solarSystemObjects['heliocentricDist'] / solarSystemObjects['topocentricDist'], axis=0 

164 ))) 

165 

166 stateVectorColumns = ['heliocentricX', 'heliocentricY', 'heliocentricZ', 'heliocentricVX', 

167 'heliocentricVY', 'heliocentricVZ', 'topocentricX', 'topocentricY', 

168 'topocentricZ', 'topocentricVX', 'topocentricVY', 'topocentricVZ'] 

169 

170 mpcorbColumns = [col for col in solarSystemObjects.columns if col[:7] == 'MPCORB_'] 

171 

172 maskedObjects = self._maskToCcdRegion( 

173 solarSystemObjects, 

174 bbox, 

175 wcs, 

176 solarSystemObjects["Err(arcsec)"].max()).copy() 

177 nSolarSystemObjects = len(maskedObjects) 

178 if nSolarSystemObjects <= 0: 

179 return self._return_empty(diaSourceCatalog, maskedObjects) 

180 

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

182 

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

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

185 diaSourceCatalog["dec"]) 

186 

187 # Create KDTree of DIA sources 

188 tree = cKDTree(vectors) 

189 

190 nFound = 0 

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

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

193 # fancier later. 

194 ssSourceData, ssObjectIds = [], [] 

195 ras, decs, residual_ras, residual_decs, dia_ids = [], [], [], [], [] 

196 diaSourceCatalog["ssObjectId"] = 0 

197 source_column = 'id' 

198 maskedObjects['associated'] = False 

199 if 'diaSourceId' in diaSourceCatalog.columns: 

200 source_column = 'diaSourceId' 

201 for ssObject in maskedObjects: 

202 index = ssObject.index 

203 ssoVect = self._radec_to_xyz(ssObject["ra"], ssObject["dec"]) 

204 # Which DIA Sources fall within r? 

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

206 if len(idx) == 1 and np.isfinite(dist[0]): 

207 nFound += 1 

208 diaSourceCatalog[idx[0]]["ssObjectId"] = ssObject["ssObjectId"] 

209 ssObjectIds.append(ssObject["ssObjectId"]) 

210 all_cols = ["phaseAngle", "heliocentricDist", 

211 "topocentricDist"] + stateVectorColumns + mpcorbColumns 

212 ssSourceData.append(list(ssObject[all_cols].values())) 

213 dia_ra = diaSourceCatalog[idx[0]]["ra"] 

214 dia_dec = diaSourceCatalog[idx[0]]["dec"] 

215 dia_id = diaSourceCatalog[idx[0]][source_column] 

216 ras.append(dia_ra) 

217 decs.append(dia_dec) 

218 dia_ids.append(dia_id) 

219 residual_ras.append(dia_ra - ssObject["ra"]) 

220 residual_decs.append(dia_dec - ssObject["dec"]) 

221 maskedObjects['associated'][index] = True 

222 else: 

223 maskedObjects['associated'][index] = False 

224 

225 self.log.info("Successfully associated %d / %d SolarSystemObjects.", nFound, nSolarSystemObjects) 

226 self.metadata['nAssociatedSsObjects'] = nFound 

227 self.metadata['nExpectedSsObjects'] = nSolarSystemObjects 

228 assocSourceMask = diaSourceCatalog["ssObjectId"] != 0 

229 unAssocObjectMask = np.logical_not(maskedObjects['associated'].value) 

230 ssSourceData = np.array(ssSourceData) 

231 ssSourceData = Table(ssSourceData, 

232 names=[ 

233 "phaseAngle", "heliocentricDist", "topocentricDist" 

234 ] + stateVectorColumns + mpcorbColumns) 

235 ssSourceData['ssObjectId'] = Column(data=ssObjectIds, dtype=int) 

236 ssSourceData["ra"] = ras 

237 ssSourceData["dec"] = decs 

238 ssSourceData["residualRa"] = residual_ras 

239 ssSourceData["residualDec"] = residual_decs 

240 ssSourceData[source_column] = dia_ids 

241 coords = SkyCoord(ra=ssSourceData['ra'].value * u.deg, dec=ssSourceData['dec'].value * u.deg) 

242 ssSourceData['galacticL'] = coords.galactic.l.deg 

243 ssSourceData['galacticB'] = coords.galactic.b.deg 

244 ssSourceData['eclipticLambda'] = coords.barycentrictrueecliptic.lon.deg 

245 ssSourceData['eclipticBeta'] = coords.barycentrictrueecliptic.lat.deg 

246 unassociatedObjects = maskedObjects[unAssocObjectMask] 

247 columns_to_drop = [ 

248 "obs_position", "obs_velocity", "obj_position", "obj_velocity", "topocentric_position", 

249 "topocentric_velocity", "obs_x_poly", "obs_y_poly", "obs_z_poly", "obj_x_poly", "obj_y_poly", 

250 "obj_z_poly", "associated" 

251 ] 

252 unassociatedObjects.remove_columns(columns_to_drop) 

253 return pipeBase.Struct( 

254 ssoAssocDiaSources=diaSourceCatalog[assocSourceMask], 

255 unAssocDiaSources=diaSourceCatalog[~assocSourceMask], 

256 nTotalSsObjects=nSolarSystemObjects, 

257 nAssociatedSsObjects=nFound, 

258 associatedSsSources=ssSourceData, 

259 unassociatedSsObjects=unassociatedObjects) 

260 

261 def _maskToCcdRegion(self, solarSystemObjects, bbox, wcs, marginArcsec): 

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

263 bounding box. 

264 

265 Parameters 

266 ---------- 

267 solarSystemObjects : `astropy.table.Table` 

268 SolarSystemObjects to mask to ``exposure``. 

269 bbox : 

270 Exposure bbox used for masking 

271 wcs : 

272 Exposure wcs used for masking 

273 marginArcsec : `float` 

274 Maximum possible matching radius to pad onto the exposure bounding 

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

276 be used. 

277 

278 Returns 

279 ------- 

280 maskedSolarSystemObjects : `astropy.table.Table` 

281 Set of SolarSystemObjects contained within the exposure bounds. 

282 """ 

283 if len(solarSystemObjects) == 0: 

284 return solarSystemObjects 

285 padding = min( 

286 int(np.ceil(marginArcsec / wcs.getPixelScale(bbox.getCenter()).asArcseconds())), 

287 self.config.maxPixelMargin) 

288 

289 return solarSystemObjects[bbox_contains_sky_coords( 

290 bbox, 

291 wcs, 

292 solarSystemObjects['ra'].value * u.degree, 

293 solarSystemObjects['dec'].value * u.degree, 

294 padding)] 

295 

296 def _radec_to_xyz(self, ras, decs): 

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

298 

299 Parameters 

300 ---------- 

301 ras : `array-like` 

302 RA coordinates of objects in degrees. 

303 decs : `array-like` 

304 DEC coordinates of objects in degrees. 

305 

306 Returns 

307 ------- 

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

309 Output unit-vectors 

310 """ 

311 ras = np.radians(ras) 

312 decs = np.radians(decs) 

313 try: 

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

315 except TypeError: 

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

317 

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

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

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

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

322 

323 return vectors 

324 

325 def _return_empty(self, diaSourceCatalog, emptySolarSystemObjects): 

326 """Return a struct with all appropriate empty values for no SSO associations. 

327 

328 Parameters 

329 ---------- 

330 diaSourceCatalog : `astropy.table.Table` 

331 Used for column names 

332 emptySolarSystemObjects : `astropy.table.Table` 

333 Used for column names. 

334 Returns 

335 ------- 

336 results : `lsst.pipe.base.Struct` 

337 Results struct with components. 

338 - ``ssoAssocDiaSources`` : Empty. (`astropy.table.Table`) 

339 - ``unAssocDiaSources`` : Input DiaSources. (`astropy.table.Table`) 

340 - ``nTotalSsObjects`` : Zero. (`int`) 

341 - ``nAssociatedSsObjects`` : Zero. 

342 - ``associatedSsSources`` : Empty. (`Astropy.table.Table`) 

343 - ``unassociatedSsObjects`` : Empty. (`Astropy.table.Table`) 

344 

345 

346 Raises 

347 ------ 

348 RuntimeError 

349 Raised if duplicate DiaObjects or duplicate DiaSources are found. 

350 """ 

351 self.log.info("No SolarSystemObjects found in detector bounding box.") 

352 return pipeBase.Struct( 

353 ssoAssocDiaSources=Table(names=diaSourceCatalog.columns), 

354 unAssocDiaSources=diaSourceCatalog, 

355 nTotalSsObjects=0, 

356 nAssociatedSsObjects=0, 

357 associatedSsSources=Table(names=['phaseAngle', 'heliocentricDist', 'topocentricDist', 

358 'heliocentricX', 'heliocentricY', 'heliocentricZ', 

359 'heliocentricVX', 'heliocentricVY', 'heliocentricVZ', 

360 'topocentricX', 'topocentricY', 'topocentricZ', 

361 'topocentricVX', 'topocentricVY', 'topocentricVZ', 

362 'residualRa', 'residualDec', 'eclipticLambda', 'eclipticBeta', 

363 'galacticL', 'galacticB', 'ssObjectId', 'diaSourceId'], 

364 dtype=[float] * 21 + [int] * 2), 

365 unassociatedSsObjects=Table(names=emptySolarSystemObjects.columns) 

366 )