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

1# This file is part of ap_association.

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.

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/>.

22"""Spatial association for Solar System Objects."""

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

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

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

38from lsst.pipe.tasks.associationUtils import obj_id_to_ss_object_id

41class SolarSystemAssociationConfig(pexConfig.Config):

42 """Config class for SolarSystemAssociationTask.

43 """

44 maxDistArcSeconds = pexConfig.Field(

45 dtype=float,

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

47 'match to a SSObject.',

48 default=1.0,

49 )

50 maxPixelMargin = pexConfig.RangeField(

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

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

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

54 "of the SolarSystemObjects in pixels.",

55 dtype=int,

56 default=100,

57 min=0,

58 )

61class SolarSystemAssociationTask(pipeBase.Task):

62 """Associate DIASources into existing SolarSystem Objects.

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

65 with known solar system objects.

66 """

67 ConfigClass = SolarSystemAssociationConfig

68 _DefaultName = "ssoAssociation"

70 @timeMethod

71 def run(self, diaSourceCatalog, ssObjects, visitInfo, bbox, wcs):

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

73 to the nearest ssoObject.

75 Parameters

76 ----------

77 diaSourceCatalog : `astropy.table.Table`

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

79 successfully associated DiaSources.

80 ssObjects : `astropy.table.Table`

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

82 of the current visit.

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

84 visitInfo of exposure used for exposure time

85 bbox : `lsst.geom.Box2I`

86 bbox of exposure used for masking

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

88 wcs of exposure used for masking

90 Returns

91 -------

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

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

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

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

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

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

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

100 - ``nAssociatedSsObjects`` : Number of SolarSystemObjects

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

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

103 """

104

105 nSolarSystemObjects = len(ssObjects)

106 if nSolarSystemObjects <= 0:

107 return self._return_empty(diaSourceCatalog, ssObjects)

108

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

110 if 'obs_x_poly' in ssObjects.columns:

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

112 ssObjects['obs_position'] = [

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

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

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

116 for row in ssObjects]

117 ssObjects['obs_velocity'] = [

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

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

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

121 for row in ssObjects]

122 ssObjects['obj_position'] = [

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

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

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

126 for row in ssObjects]

127 ssObjects['obj_velocity'] = [

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

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

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

131 for row in ssObjects]

132 vector = np.vstack(ssObjects['obj_position'].value - ssObjects['obs_position'].value)

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

134 ssObjects['ra'] = ras

135 ssObjects['dec'] = decs

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

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

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

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

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

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

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

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

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

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

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

147 )

148 ssObjects['topocentricX'], ssObjects['topocentricY'], ssObjects['topocentricZ'] = (

149 np.array(list(ssObjects['topocentric_position'].value)).T

150 )

151 ssObjects['topocentricVX'], ssObjects['topocentricVY'], ssObjects['topocentricVZ'] = (

152 np.array(list(ssObjects['topocentric_velocity'].value)).T

153 )

154 ssObjects['heliocentricVX'], ssObjects['heliocentricVY'], \

155 ssObjects['heliocentricVZ'] = np.array(list(ssObjects['obj_velocity'].value)).T

156 ssObjects['heliocentricDist'] = np.linalg.norm(ssObjects['obj_position'], axis=1)

157 ssObjects['topocentricDist'] = np.linalg.norm(ssObjects['topocentric_position'], axis=1)

158 ssObjects['phaseAngle'] = np.degrees(np.arccos(np.sum(

159 ssObjects['obj_position'].T * ssObjects['topocentric_position'].T

160 / ssObjects['heliocentricDist'] / ssObjects['topocentricDist'], axis=0

161 )))

162 marginArcsec = ssObjects["Err(arcsec)"].max()

163

164 columns_to_drop = [

165 "obs_position", "obs_velocity", "obj_position", "obj_velocity", "topocentric_position",

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

167 "obj_z_poly", "associated"

168 ]

169

170 else:

171 ssObjects.rename_columns(

172 ['RA_deg', 'Dec_deg', 'phase_deg', 'Range_LTC_km', 'Obj_Sun_x_LTC_km', 'Obj_Sun_y_LTC_km',

173 'Obj_Sun_z_LTC_km', 'Obj_Sun_vx_LTC_km_s', 'Obj_Sun_vy_LTC_km_s', 'Obj_Sun_vz_LTC_km_s'],

174 ['ra', 'dec', 'phaseAngle', 'topocentricDist', 'heliocentricX', 'heliocentricY',

175 'heliocentricZ', 'heliocentricVX', 'heliocentricVY', 'heliocentricVZ'])

176 ssObjects['ssObjectId'] = [obj_id_to_ss_object_id(v) for v in ssObjects['ObjID']]

177 ssObjects['heliocentricDist'] = (

178 np.sqrt(ssObjects['heliocentricX']**2 + ssObjects['heliocentricY']**2

179 + ssObjects['heliocentricZ']**2)

180 )

181 for substring1, substring2 in [('X', 'x_km'), ('Y', 'y_km'), ('Z', 'z_km'),

182 ('VX', 'vx_km_s'), ('VY', 'vy_km_s'), ('VZ', 'vz_km_s')]:

183 topoName = 'topocentric' + substring1

184 helioName = 'heliocentric' + substring1

185 obsName = 'Obs_Sun_' + substring2

186 ssObjects[topoName] = ssObjects[helioName] - ssObjects[obsName]

187

188 marginArcsec = 1.0 # TODO: justify

189

190 columns_to_drop = ['FieldID', 'fieldMJD_TAI', 'fieldJD_TDB',

191 'RangeRate_LTC_km_s', 'RARateCosDec_deg_day',

192 'DecRate_deg_day', 'Obs_Sun_x_km', 'Obs_Sun_y_km', 'Obs_Sun_z_km',

193 'Obs_Sun_vx_km_s', 'Obs_Sun_vy_km_s', 'Obs_Sun_vz_km_s',

194 '__index_level_0__']

195

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

197 'heliocentricVY', 'heliocentricVZ', 'topocentricX', 'topocentricY',

198 'topocentricZ', 'topocentricVX', 'topocentricVY', 'topocentricVZ']

199

200 mpcorbColumns = [col for col in ssObjects.columns if col[:7] == 'MPCORB_']

201

202 maskedObjects = self._maskToCcdRegion(

203 ssObjects,

204 bbox,

205 wcs,

206 marginArcsec).copy()

207 nSolarSystemObjects = len(maskedObjects)

208 if nSolarSystemObjects <= 0:

209 return self._return_empty(diaSourceCatalog, maskedObjects)

210

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

212

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

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

215 diaSourceCatalog["dec"])

216

217 # Create KDTree of DIA sources

218 tree = cKDTree(vectors)

219

220 nFound = 0

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

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

223 # fancier later.

224 ssSourceData, ssObjectIds = [], []

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

226 diaSourceCatalog["ssObjectId"] = 0

227 source_column = 'id'

228 maskedObjects['associated'] = False

229 if 'diaSourceId' in diaSourceCatalog.columns:

230 source_column = 'diaSourceId'

231 for ssObject in maskedObjects:

232 index = ssObject.index

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

234 # Which DIA Sources fall within r?

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

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

237 nFound += 1

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

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

240 all_cols = ["phaseAngle", "heliocentricDist",

241 "topocentricDist"] + stateVectorColumns + mpcorbColumns

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

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

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

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

246 ras.append(dia_ra)

247 decs.append(dia_dec)

248 dia_ids.append(dia_id)

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

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

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

252 else:

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

254

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

256 self.metadata['nAssociatedSsObjects'] = nFound

257 self.metadata['nExpectedSsObjects'] = nSolarSystemObjects

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

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

260 ssSourceData = np.array(ssSourceData)

261 ssSourceData = Table(ssSourceData,

262 names=[

263 "phaseAngle", "heliocentricDist", "topocentricDist"

264 ] + stateVectorColumns + mpcorbColumns)

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

266 ssSourceData["ra"] = ras

267 ssSourceData["dec"] = decs

268 ssSourceData["residualRa"] = residual_ras

269 ssSourceData["residualDec"] = residual_decs

270 ssSourceData[source_column] = dia_ids

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

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

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

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

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

276 unassociatedObjects = maskedObjects[unAssocObjectMask]

277 unassociatedObjects.remove_columns(columns_to_drop)

278 return pipeBase.Struct(

279 ssoAssocDiaSources=diaSourceCatalog[assocSourceMask],

280 unAssocDiaSources=diaSourceCatalog[~assocSourceMask],

281 nTotalSsObjects=nSolarSystemObjects,

282 nAssociatedSsObjects=nFound,

283 associatedSsSources=ssSourceData,

284 unassociatedSsObjects=unassociatedObjects)

285

286 def _maskToCcdRegion(self, ssObjects, bbox, wcs, marginArcsec):

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

288 bounding box.

289

290 Parameters

291 ----------

292 ssObjects : `astropy.table.Table`

293 SolarSystemObjects to mask to ``exposure``.

294 bbox :

295 Exposure bbox used for masking

296 wcs :

297 Exposure wcs used for masking

298 marginArcsec : `float`

299 Maximum possible matching radius to pad onto the exposure bounding

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

301 be used.

302

303 Returns

304 -------

305 maskedSolarSystemObjects : `astropy.table.Table`

306 Set of SolarSystemObjects contained within the exposure bounds.

307 """

308 if len(ssObjects) == 0:

309 return ssObjects

310 padding = min(

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

312 self.config.maxPixelMargin)

313

314 return ssObjects[bbox_contains_sky_coords(

315 bbox,

316 wcs,

317 ssObjects['ra'].value * u.degree,

318 ssObjects['dec'].value * u.degree,

319 padding)]

320

321 def _radec_to_xyz(self, ras, decs):

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

323

324 Parameters

325 ----------

326 ras : `array-like`

327 RA coordinates of objects in degrees.

328 decs : `array-like`

329 DEC coordinates of objects in degrees.

330

331 Returns

332 -------

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

334 Output unit-vectors

335 """

336 ras = np.radians(ras)

337 decs = np.radians(decs)

338 try:

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

340 except TypeError:

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

342

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

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

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

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

347

348 return vectors

349

350 def _return_empty(self, diaSourceCatalog, emptySolarSystemObjects):

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

352

353 Parameters

354 ----------

355 diaSourceCatalog : `astropy.table.Table`

356 Used for column names

357 emptySolarSystemObjects : `astropy.table.Table`

358 Used for column names.

359 Returns

360 -------

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

362 Results struct with components.

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

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

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

366 - ``nAssociatedSsObjects`` : Zero.

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

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

369

370

371 Raises

372 ------

373 RuntimeError

374 Raised if duplicate DiaObjects or duplicate DiaSources are found.

375 """

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

377 return pipeBase.Struct(

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

379 unAssocDiaSources=diaSourceCatalog,

380 nTotalSsObjects=0,

381 nAssociatedSsObjects=0,

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

383 'heliocentricX', 'heliocentricY', 'heliocentricZ',

384 'heliocentricVX', 'heliocentricVY', 'heliocentricVZ',

385 'topocentricX', 'topocentricY', 'topocentricZ',

386 'topocentricVX', 'topocentricVY', 'topocentricVZ',

387 'residualRa', 'residualDec', 'eclipticLambda', 'eclipticBeta',

388 'galacticL', 'galacticB', 'ssObjectId', 'diaSourceId'],

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

390 unassociatedSsObjects=Table(names=emptySolarSystemObjects.columns)

391 )