Coverage for python/lsst/ap/association/association.py: 98%
139 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-01 02:13 -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/>.
22"""A simple implementation of source association task for ap_verify.
23"""
25__all__ = ["AssociationConfig", "AssociationTask"]
27import itertools
29import numpy as np
30import pandas as pd
31from scipy.sparse import csr_matrix
32from scipy.sparse.csgraph import min_weight_full_bipartite_matching
33from scipy.spatial import cKDTree
35import lsst.geom as geom
36import lsst.pex.config as pexConfig
37import lsst.pipe.base as pipeBase
38from lsst.pipe.tasks.schemaUtils import column_dtype
39from lsst.utils.timer import timeMethod
41# Enforce an error for unsafe column/array value setting in pandas.
42pd.options.mode.chained_assignment = 'raise'
45class AssociationConfig(pexConfig.Config):
46 """Config class for AssociationTask.
47 """
49 maxDistArcSeconds = pexConfig.Field(
50 dtype=float,
51 doc="Maximum distance in arcseconds for a DIASource to be matched "
52 "to a DIAObject. This is the sole association radius: every "
53 "DIAObject within this distance is a match candidate, ranked by "
54 "position chi^2 when uncertainties are available and by angular "
55 "distance otherwise.",
56 default=1.0,
57 )
58 sigmaFloorArcSeconds = pexConfig.Field(
59 dtype=float,
60 doc="Floor on the per-axis position uncertainty (arcsec) used when "
61 "computing chi^2.",
62 default=0.05,
63 )
64 fallbackSigmaArcSeconds = pexConfig.Field(
65 dtype=float,
66 doc="Per-axis position uncertainty (arcsec) substituted for "
67 "raErr/decErr values that are missing, non-finite, or "
68 "non-positive when computing chi^2. Should reflect a plausible "
69 "size for an unmeasured per-axis uncertainty (e.g., one LSSTCam "
70 "pixel ~= 0.2 arcsec). Has no effect on rows that carry valid "
71 "uncertainty values.",
72 default=0.2,
73 )
76class AssociationTask(pipeBase.Task):
77 """Associate DIAOSources into existing DIAObjects.
79 This task performs the association of detected DIASources in a visit
80 with the previous DIAObjects detected over time. It also creates new
81 DIAObjects out of DIASources that cannot be associated with previously
82 detected DIAObjects.
83 """
85 ConfigClass = AssociationConfig
86 _DefaultName = "association"
88 @timeMethod
89 def run(self,
90 diaSources,
91 diaObjects,
92 schema=None):
93 """Associate the new DiaSources with existing DiaObjects.
95 Parameters
96 ----------
97 diaSources : `pandas.DataFrame`
98 New DIASources to be associated with existing DIAObjects.
99 diaObjects : `pandas.DataFrame`
100 Existing diaObjects from the Apdb.
101 schema : `dict` [`str`, `felis.datamodel.Schema`] or `None`, optional
102 Dictionary of Schemas from ``sdm_schemas`` containing the table
103 definition to use. If `None`, dtypes for new columns are guessed
104 from the input tables.
106 Returns
107 -------
108 result : `lsst.pipe.base.Struct`
109 Results struct with components.
111 - ``matchedDiaSources`` : DiaSources that were matched. Matched
112 Sources have their diaObjectId updated and set to the id of the
113 diaObject they were matched to. (`pandas.DataFrame`)
114 - ``unAssocDiaSources`` : DiaSources that were not matched.
115 Unassociated sources have their diaObject set to 0 as they
116 were not associated with any existing DiaObjects.
117 (`pandas.DataFrame`)
118 - ``nUpdatedDiaObjects`` : Number of DiaObjects that were
119 matched to new DiaSources. (`int`)
120 - ``nUnassociatedDiaObjects`` : Number of DiaObjects that were
121 not matched a new DiaSource. (`int`)
122 """
123 diaSources = self.check_dia_source_radec(diaSources)
125 if len(diaObjects) == 0:
126 return pipeBase.Struct(
127 matchedDiaSources=pd.DataFrame(columns=diaSources.columns),
128 unAssocDiaSources=diaSources,
129 nUpdatedDiaObjects=0,
130 nUnassociatedDiaObjects=0)
132 matchResult = self.associate_sources(diaObjects, diaSources, schema)
134 mask = matchResult.diaSources["diaObjectId"] != 0
136 return pipeBase.Struct(
137 matchedDiaSources=matchResult.diaSources[mask].reset_index(drop=True),
138 unAssocDiaSources=matchResult.diaSources[~mask].reset_index(drop=True),
139 nUpdatedDiaObjects=matchResult.nUpdatedDiaObjects,
140 nUnassociatedDiaObjects=matchResult.nUnassociatedDiaObjects)
142 def check_dia_source_radec(self, dia_sources):
143 """Check that all DiaSources have non-NaN values for RA/DEC.
145 If one or more DiaSources are found to have NaN values, throw a
146 warning to the log with the ids of the offending sources. Drop them
147 from the table.
149 Parameters
150 ----------
151 dia_sources : `pandas.DataFrame`
152 Input DiaSources to check for NaN values.
154 Returns
155 -------
156 trimmed_sources : `pandas.DataFrame`
157 DataFrame of DiaSources trimmed of all entries with NaN values for
158 RA/DEC.
159 """
160 nan_mask = dia_sources["ra"].isnull() | dia_sources["dec"].isnull()
161 if nan_mask.any():
162 nan_ids = dia_sources.loc[nan_mask, "diaSourceId"]
163 for nan_id in nan_ids:
164 self.log.warning(
165 "DiaSource %i has NaN value for RA/DEC, "
166 "dropping from association.", nan_id)
167 dia_sources = dia_sources[~nan_mask]
168 return dia_sources
170 @timeMethod
171 def associate_sources(self, dia_objects, dia_sources, schema=None):
172 """Associate the input DIASources with the catalog of DIAObjects.
174 DiaObject DataFrame must be indexed on ``diaObjectId``.
176 Parameters
177 ----------
178 dia_objects : `pandas.DataFrame`
179 Catalog of DIAObjects to attempt to associate the input
180 DIASources into.
181 dia_sources : `pandas.DataFrame`
182 DIASources to associate into the DIAObjectCollection.
183 schema : `dict` [`str`, `felis.datamodel.Schema`] or `None`, optional
184 Dictionary of Schemas from ``sdm_schemas`` containing the table
185 definition to use.
187 Returns
188 -------
189 result : `lsst.pipe.base.Struct`
190 Results struct with components.
192 - ``diaSources`` : Full set of diaSources both matched and not.
193 (`pandas.DataFrame`)
194 - ``nUpdatedDiaObjects`` : Number of DiaObjects that were
195 associated. (`int`)
196 - ``nUnassociatedDiaObjects`` : Number of DiaObjects that were
197 not matched a new DiaSource. (`int`)
198 """
199 scores = self.score(
200 dia_objects, dia_sources,
201 self.config.maxDistArcSeconds * geom.arcseconds)
202 match_result = self.match(dia_objects, dia_sources, scores, schema)
204 return match_result
206 @timeMethod
207 def score(self, dia_objects, dia_sources, max_dist):
208 """Build the candidate (DIASource, DIAObject) match table and
209 score every pair.
211 For each DIASource, all DIAObjects within ``max_dist`` are retrieved
212 from a kd-tree on unit vectors. Each candidate pair is then scored:
214 - If both inputs carry usable ``raErr``/``decErr`` columns, the
215 score is the 2D Gaussian negative log-likelihood of the
216 position residual (``0.5 * chi^2 + 0.5 * ln(var_ra * var_dec)``),
217 so the match prefers the most likely object, not merely the
218 nearest or the lowest-chi^2 one (see `_position_nll`).
219 - Otherwise, the distance (in radians) is used as the score.
221 No candidates are dropped by the score itself: every pair within
222 ``max_dist`` is retained, and the score is used only to rank them.
224 ``raErr`` and ``decErr`` are taken to follow the LSST DPDD
225 convention: each is the marginal uncertainty of the catalog
226 coordinate itself in degrees (no cos(dec) factor folded into
227 ``raErr``). Under that convention the cos(dec) factor cancels
228 between residual and uncertainty, and chi^2 reduces to
229 ``dRA^2 / sum(raErr^2) + dDec^2 / sum(decErr^2)``.
231 ``max_dist`` is both the candidate pre-filter and the
232 association radius: every pair within it is retained, and the
233 score is used only to rank candidates in the downstream match.
235 Parameters
236 ----------
237 dia_objects, dia_sources : `pandas.DataFrame`
238 Must contain ``ra`` and ``dec``; ``raErr`` and ``decErr`` are
239 used when present.
240 max_dist : `lsst.geom.Angle`
241 Hard angular upper bound on candidate pairs.
243 Returns
244 -------
245 result : `lsst.pipe.base.Struct`
246 Flat candidate-pair table:
248 - ``src_idx`` : `numpy.ndarray` of `int`
249 Positional source index for each surviving pair.
250 - ``obj_idx`` : `numpy.ndarray` of `int`
251 Positional object index for each surviving pair.
252 - ``scores`` : `numpy.ndarray` of `float`
253 Cost of each pair (position negative log-likelihood if
254 uncertainty-based, chord distance in radians otherwise).
255 Lower is better; NLL values may be negative.
256 - ``unmatched_cost`` : `float`
257 Cost to assign to the synthetic 'no-match' alternative
258 in the linear-assignment match — set so that any
259 surviving real candidate is preferred.
260 """
261 n_src = len(dia_sources)
262 n_obj = len(dia_objects)
263 empty_int = np.empty(0, dtype=np.int64)
264 empty_float = np.empty(0, dtype=np.float64)
265 max_dist_rad = max_dist.asRadians()
266 # Used as the no-match cost in distance mode; always strictly
267 # above any real candidate (which the kd-tree caps at max_dist_rad).
268 chord_unmatched_cost = max_dist_rad * 1.01 + 1e-300
270 if n_obj == 0 or n_src == 0: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 return pipeBase.Struct(
272 src_idx=empty_int,
273 obj_idx=empty_int,
274 scores=empty_float,
275 unmatched_cost=chord_unmatched_cost)
277 spatial_tree = self._make_spatial_tree(dia_objects)
278 src_vectors = self._radec_to_xyz(dia_sources)
280 candidate_lists = spatial_tree.query_ball_point(
281 src_vectors, r=max_dist_rad)
282 counts = np.fromiter(
283 (len(c) for c in candidate_lists), dtype=np.int64, count=n_src)
284 n_pairs = int(counts.sum())
285 if n_pairs == 0:
286 return pipeBase.Struct(
287 src_idx=empty_int,
288 obj_idx=empty_int,
289 scores=empty_float,
290 unmatched_cost=chord_unmatched_cost)
292 src_idx = np.repeat(np.arange(n_src, dtype=np.int64), counts)
293 obj_idx = np.fromiter(
294 itertools.chain.from_iterable(candidate_lists),
295 dtype=np.int64, count=n_pairs)
297 if (self._has_position_errors(dia_sources)
298 and self._has_position_errors(dia_objects)):
299 scores = self._position_nll(dia_sources, dia_objects, src_idx, obj_idx)
300 # ``max_dist`` is the sole association gate: every candidate
301 # pair is kept and the NLL only ranks them. Price the
302 # 'no-match' alternative just above the worst candidate so a
303 # real match is always preferred. A diaSource will only not be
304 # associated if there are more diaSources than diaObjects.
305 unmatchedCostDelta = 1.0
306 unmatched_cost = (float(scores.max()) + unmatchedCostDelta)
307 else:
308 obj_vectors = self._radec_to_xyz(dia_objects)
309 diffs = src_vectors[src_idx] - obj_vectors[obj_idx]
310 scores = np.linalg.norm(diffs, axis=1)
311 unmatched_cost = chord_unmatched_cost
313 return pipeBase.Struct(
314 src_idx=src_idx,
315 obj_idx=obj_idx,
316 scores=scores,
317 unmatched_cost=unmatched_cost)
319 @staticmethod
320 def _has_position_errors(catalog):
321 """Return True iff ``catalog`` carries ``raErr`` and ``decErr``
322 columns with at least one finite, positive value in each.
323 """
324 if "raErr" not in catalog.columns or "decErr" not in catalog.columns:
325 return False
326 raErr = catalog["raErr"].to_numpy()
327 decErr = catalog["decErr"].to_numpy()
328 return (bool(np.any(np.isfinite(raErr) & (raErr > 0.0)))
329 and bool(np.any(np.isfinite(decErr) & (decErr > 0.0))))
331 def _position_nll(self, dia_sources, dia_objects, src_idx, obj_idx):
332 """Return the position Gaussian negative log-likelihood (NLL) for
333 paired DIASources/DIAObjects, used as the match cost.
335 Under the hypothesis that the DIASource and DIAObject are the same
336 astrophysical object, each per-axis residual is a zero-mean
337 Gaussian whose variance is the sum of the two catalogs' squared
338 per-axis uncertainties. The cost is the negative log-likelihood of
339 the observed residual, dropping the constant ``ln(2*pi)`` term:
341 NLL = 0.5 * (dra^2 / var_ra + ddec^2 / var_dec)
342 + 0.5 * ln(var_ra * var_dec).
344 The first term is half the 2-DOF position chi^2. The second is the
345 Gaussian normalization, which penalises poorly-localised
346 candidates: a bare chi^2 cost omits it, and since a larger
347 variance deflates chi^2 for a fixed separation, a bare chi^2 would
348 bias the match toward the more poorly-localised object.
350 Non-finite or non-positive per-row uncertainties are replaced
351 with ``self.config.fallbackSigmaArcSeconds``. The combined per-axis
352 variance is then floored at ``self.config.sigmaFloorArcSeconds``
353 to guard against pathologically small reported errors.
355 Parameters
356 ----------
357 dia_sources, dia_objects : `pandas.DataFrame`
358 Catalogs containing ``ra``, ``dec``, ``raErr``, ``decErr``
359 (all in degrees).
360 src_idx, obj_idx : `numpy.ndarray` of `int`
361 Paired positional indices; ``src_idx[k]`` is matched against
362 ``obj_idx[k]``.
364 Returns
365 -------
366 nll : `numpy.ndarray` of `float`
367 Position negative log-likelihood, one value per pair. Lower is
368 a better (more likely) match; values may be negative.
369 """
370 sigma_floor_sq_deg = (self.config.sigmaFloorArcSeconds / 3600.0) ** 2
371 fallback_sq_deg = (self.config.fallbackSigmaArcSeconds / 3600.0) ** 2
373 def err_sq(catalog, col, idx):
374 arr = catalog[col].to_numpy()[idx]
375 sq = np.square(arr)
376 return np.where(np.isfinite(sq) & (arr > 0.0),
377 sq, fallback_sq_deg)
379 src_ra = dia_sources["ra"].to_numpy()[src_idx]
380 src_dec = dia_sources["dec"].to_numpy()[src_idx]
381 obj_ra = dia_objects["ra"].to_numpy()[obj_idx]
382 obj_dec = dia_objects["dec"].to_numpy()[obj_idx]
384 # Make sure that the RA difference is never greater than +/-180 in
385 # either direction.
386 dra = ((src_ra - obj_ra) + 180.0) % 360.0 - 180.0
387 ddec = src_dec - obj_dec
389 var_ra = (err_sq(dia_sources, "raErr", src_idx)
390 + err_sq(dia_objects, "raErr", obj_idx))
391 var_dec = (err_sq(dia_sources, "decErr", src_idx)
392 + err_sq(dia_objects, "decErr", obj_idx))
393 var_ra = np.maximum(var_ra, sigma_floor_sq_deg)
394 var_dec = np.maximum(var_dec, sigma_floor_sq_deg)
396 chi2 = dra**2/var_ra + ddec**2/var_dec
397 return 0.5*chi2 + 0.5*np.log(var_ra*var_dec)
399 def _make_spatial_tree(self, dia_objects):
400 """Create a searchable kd-tree the input dia_object positions.
402 Parameters
403 ----------
404 dia_objects : `pandas.DataFrame`
405 A catalog of DIAObjects to create the tree from.
407 Returns
408 -------
409 kd_tree : `scipy.spatical.cKDTree`
410 Searchable kd-tree created from the positions of the DIAObjects.
411 """
412 vectors = self._radec_to_xyz(dia_objects)
413 return cKDTree(vectors)
415 def _radec_to_xyz(self, catalog):
416 """Convert input ra/dec coordinates to spherical unit-vectors.
418 Parameters
419 ----------
420 catalog : `pandas.DataFrame`
421 Catalog to produce spherical unit-vector from.
423 Returns
424 -------
425 vectors : `numpy.ndarray`, (N, 3)
426 Output unit-vectors
427 """
428 ras = np.radians(catalog["ra"])
429 decs = np.radians(catalog["dec"])
430 vectors = np.empty((len(ras), 3))
432 sin_dec = np.sin(np.pi / 2 - decs)
433 vectors[:, 0] = sin_dec * np.cos(ras)
434 vectors[:, 1] = sin_dec * np.sin(ras)
435 vectors[:, 2] = np.cos(np.pi / 2 - decs)
437 return vectors
439 @timeMethod
440 def match(self, dia_objects, dia_sources, score_struct, schema=None):
441 """Solve a min-cost bipartite matching between sources and objects.
443 Each DIASource is given a synthetic 'no-match' alternative (a
444 per-source 'ghost' column) carrying ``score_struct.unmatched_cost``.
445 A min-weight full bipartite matching is then solved on the sparse
446 cost matrix made from real candidate pairs and ghost edges.
447 Sources matched to a ghost are reported as unassociated; sources
448 matched to a real object inherit that object's ``diaObjectId``.
450 When two sources compete for the same object, the source with a
451 strictly worse alternative gets that object, and the other source falls
452 back to its second-best candidate rather than creating a new DIAObject.
454 Parameters
455 ----------
456 dia_objects, dia_sources : `pandas.DataFrame`
457 Must contain ``ra`` and ``dec``; ``raErr`` and ``decErr`` are
458 used when present.
459 score_struct : `lsst.pipe.base.Struct`
460 Output of `score`: ``src_idx``, ``obj_idx``, ``scores``,
461 ``unmatched_cost``.
462 schema : `dict` [`str`, `felis.datamodel.Schema`] or `None`, optional
463 Dictionary of Schemas from ``sdm_schemas`` containing the table
464 definition to use.
466 Returns
467 -------
468 result : `lsst.pipe.base.Struct`
470 - ``diaSources`` : input source table with ``diaObjectId``
471 populated (0 for unmatched). (`pandas.DataFrame`)
472 - ``nUpdatedDiaObjects`` : number of DIAObjects matched to a
473 new DIASource. (`int`)
474 - ``nUnassociatedDiaObjects`` : number of preloaded DIAObjects
475 with no matching DIASource. (`int`)
476 """
477 n_src = len(dia_sources)
478 n_obj = len(dia_objects)
479 if not pd.api.types.is_integer_dtype(dia_sources["diaObjectId"]):
480 raise ValueError(f"diaSource column diaObjectId must be an integer, "
481 f"got {dia_sources['diaObjectId'].dtype} instead")
482 # Set the diaObjectId dtype from the schema if available. Otherwise
483 # fall back on the dtype of the incoming diaObjectId column.
484 # If the dtype is incompatible with the final schema (uint vs int),
485 # then pandas will silently promote the column to a float.
486 if schema is not None and schema.get("DiaSource") is not None:
487 obj_id_col = next(
488 c for c in schema["DiaSource"].columns if c.name == "diaObjectId"
489 )
490 obj_id_dtype = column_dtype(obj_id_col.datatype, nullable=obj_id_col.nullable)
491 else:
492 obj_id_dtype = dia_sources["diaObjectId"].dtype
493 # Allocate via pandas so pandas-extension dtypes (e.g., "Int64")
494 # are supported alongside numpy dtypes.
495 associated_dia_object_ids = pd.array([0]*n_src, dtype=obj_id_dtype)
496 n_matched = 0
498 if n_src > 0: 498 ↛ 535line 498 didn't jump to line 535 because the condition on line 498 was always true
499 # Per-source ghost edges to columns [n_obj, n_obj + n_src).
500 ghost_src = np.arange(n_src, dtype=np.int64)
501 ghost_obj = n_obj + ghost_src
502 ghost_cost = np.full(
503 n_src, score_struct.unmatched_cost, dtype=np.float64)
505 real_weights = score_struct.scores.astype(np.float64, copy=False)
506 rows = np.concatenate(
507 [score_struct.src_idx.astype(np.int64, copy=False),
508 ghost_src])
509 cols = np.concatenate(
510 [score_struct.obj_idx.astype(np.int64, copy=False),
511 ghost_obj])
512 weights = np.concatenate([real_weights, ghost_cost])
514 # scipy's min_weight_full_bipartite_matching treats an explicit
515 # zero as a missing edge, and the NLL cost can be negative or
516 # zero. A full matching uses exactly one edge per source, so
517 # shifting every stored weight by a constant leaves the optimal
518 # matching unchanged; rescale so the smallest weight is strictly
519 # positive and no real edge is dropped.
520 weights = weights - weights.min() + 1.0
522 biadj = csr_matrix(
523 (weights, (rows, cols)),
524 shape=(n_src, n_obj + n_src))
525 match_src, match_dst = min_weight_full_bipartite_matching(biadj)
527 is_real = match_dst < n_obj
528 matched_src = match_src[is_real]
529 matched_obj = match_dst[is_real]
530 n_matched = int(matched_src.size)
531 if n_matched > 0:
532 associated_dia_object_ids[matched_src] = (
533 dia_objects.index.to_numpy()[matched_obj])
535 dia_sources = dia_sources.copy()
536 dia_sources["diaObjectId"] = associated_dia_object_ids
538 return pipeBase.Struct(
539 diaSources=dia_sources,
540 nUpdatedDiaObjects=n_matched,
541 nUnassociatedDiaObjects=int(n_obj - n_matched))