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

Hot-keys 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
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/>.
22"""Spatial association for Solar System Objects."""
24__all__ = ["SolarSystemAssociationConfig", "SolarSystemAssociationTask"]
26import numpy as np
27import pandas as pd
28from scipy.spatial import cKDTree
30import lsst.geom as geom
31import lsst.pex.config as pexConfig
32import lsst.pipe.base as pipeBase
35class SolarSystemAssociationConfig(pexConfig.Config):
36 """Config class for SolarSystemAssociationTask.
37 """
38 maxDistArcSeconds = pexConfig.Field(
39 dtype=float,
40 doc='Maximum distance in arcseconds to test for a DIASource to be a '
41 'match to a SSObject.',
42 default=2.0,
43 )
44 maxPixelMargin = pexConfig.RangeField(
45 doc="Maximum padding to add to the ccd bounding box before masking "
46 "SolarSystem objects to the ccd footprint. The bounding box will "
47 "be padded by the minimum of this number or the max uncertainty "
48 "of the SolarSystemObjects in pixels.",
49 dtype=int,
50 default=100,
51 min=0,
52 )
55class SolarSystemAssociationTask(pipeBase.Task):
56 """Associate DIASources into existing SolarSystem Objects.
58 This task performs the association of detected DIASources in a visit
59 with known solar system objects.
60 """
61 ConfigClass = SolarSystemAssociationConfig
62 _DefaultName = "ssoAssociation"
64 @pipeBase.timeMethod
65 def run(self, diaSourceCatalog, solarSystemObjects, exposure):
66 """Create a searchable tree of unassociated DiaSources and match
67 to the nearest ssoObject.
69 Parameters
70 ----------
71 diaSourceCatalog : `pandas.DataFrame`
72 Catalog of DiaSources. Modified in place to add ssObjectId to
73 successfully associated DiaSources.
74 solarSystemObjects : `pandas.DataFrame`
75 Set of solar system objects that should be within the footprint
76 of the current visit.
77 exposure : `lsst.afw.image.ExposureF`
78 Exposure where the DiaSources in ``diaSourceCatalog`` were
79 detected in.
81 Returns
82 -------
83 resultsStruct : `lsst.pipe.base.Struct`
85 - ``ssoAssocDiaSources`` : DiaSources that were associated with
86 solar system objects in this visit. (`pandas.DataFrame`)
87 - ``unAssocDiaSources`` : Set of DiaSources that were not
88 associated with any solar system object. (`pandas.DataFrame`)
89 - ``nTotalSsObjects`` : Total number of SolarSystemObjects
90 contained in the CCD footprint. (`int`)
91 - ``nAssociatedSsObjects`` : Number of SolarSystemObjects
92 that were associated with DiaSources.
93 """
94 maskedObjects = self._maskToCcdRegion(
95 solarSystemObjects,
96 exposure,
97 solarSystemObjects["Err(arcsec)"].max())
98 nSolarSystemObjects = len(maskedObjects)
99 if nSolarSystemObjects <= 0:
100 self.log.info("No SolarSystemObjects found in detector bounding "
101 "box.")
102 return pipeBase.Struct(
103 ssoAssocDiaSources=pd.DataFrame(columns=diaSourceCatalog.columns),
104 unAssocDiaSources=diaSourceCatalog,
105 nTotalSsObjects=0,
106 nAssociatedSsObjects=0)
107 self.log.info(f"Attempting to associate {nSolarSystemObjects}...")
108 maxRadius = np.deg2rad(self.config.maxDistArcSeconds / 3600)
110 # Transform DIA RADEC coordinates to unit sphere xyz for tree building.
111 vectors = self._radec_to_xyz(diaSourceCatalog["ra"],
112 diaSourceCatalog["decl"])
114 # Create KDTree of DIA sources
115 tree = cKDTree(vectors)
117 nFound = 0
118 # Query the KDtree for DIA nearest neighbors to SSOs. Currently only
119 # picks the DiaSource with the shortest distance. We can do something
120 # fancier later.
121 for index, ssObject in maskedObjects.iterrows():
123 ssoVect = self._radec_to_xyz(ssObject["ra"], ssObject["decl"])
125 # Which DIA Sources fall within r?
126 dist, idx = tree.query(ssoVect, distance_upper_bound=maxRadius)
127 if np.isfinite(dist[0]):
128 nFound += 1
129 diaSourceCatalog.loc[idx[0], "ssObjectId"] = ssObject["ssObjectId"]
131 self.log.info(f"Successfully associated {nFound} SolarSystemObjects.")
132 assocMask = diaSourceCatalog["ssObjectId"] != 0
133 return pipeBase.Struct(
134 ssoAssocDiaSources=diaSourceCatalog[assocMask].reset_index(drop=True),
135 unAssocDiaSources=diaSourceCatalog[~assocMask].reset_index(drop=True),
136 nTotalSsObjects=nSolarSystemObjects,
137 nAssociatedSsObjects=nFound)
139 def _maskToCcdRegion(self, solarSystemObjects, exposure, marginArcsec):
140 """Mask the input SolarSystemObjects to only those in the exposure
141 bounding box.
143 Parameters
144 ----------
145 solarSystemObjects : `pandas.DataFrame`
146 SolarSystemObjects to mask to ``exposure``.
147 exposure : `lsst.afw.image.ExposureF`
148 Exposure to mask to.
149 marginArcsec : `float`
150 Maximum possible matching radius to pad onto the exposure bounding
151 box. If greater than ``maxPixelMargin``, ``maxPixelMargin`` will
152 be used.
154 Returns
155 -------
156 maskedSolarSystemObjects : `pandas.DataFrame`
157 Set of SolarSystemObjects contained within the exposure bounds.
158 """
159 wcs = exposure.getWcs()
160 padding = min(
161 int(np.ceil(wcs.getPixelScale().asArcseconds()*marginArcsec)),
162 self.config.maxPixelMargin)
163 bbox = geom.Box2D(exposure.getBBox())
164 bbox.grow(padding)
166 mapping = wcs.getTransform().getMapping()
167 x, y = mapping.applyInverse(
168 np.radians(solarSystemObjects[['ra', 'decl']].T.to_numpy()))
169 return solarSystemObjects[bbox.contains(x, y)]
171 def _radec_to_xyz(self, ras, decs):
172 """Convert input ra/dec coordinates to spherical unit-vectors.
174 Parameters
175 ----------
176 ras : `array-like`
177 RA coordinates of objects in degrees.
178 decs : `array-like`
179 DEC coordinates of objects in degrees.
181 Returns
182 -------
183 vectors : `numpy.ndarray`, (N, 3)
184 Output unit-vectors
185 """
186 ras = np.radians(ras)
187 decs = np.radians(decs)
188 try:
189 vectors = np.empty((len(ras), 3))
190 except TypeError:
191 vectors = np.empty((1, 3))
193 sin_dec = np.sin(np.pi / 2 - decs)
194 vectors[:, 0] = sin_dec * np.cos(ras)
195 vectors[:, 1] = sin_dec * np.sin(ras)
196 vectors[:, 2] = np.cos(np.pi / 2 - decs)
198 return vectors