diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py index 0b5b0b13..4ca0329e 100644 --- a/python/lsst/ap/association/association.py +++ b/python/lsst/ap/association/association.py @@ -24,13 +24,18 @@ __all__ = ["AssociationConfig", "AssociationTask"] +import itertools + import numpy as np import pandas as pd +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import min_weight_full_bipartite_matching from scipy.spatial import cKDTree import lsst.geom as geom import lsst.pex.config as pexConfig import lsst.pipe.base as pipeBase +from lsst.pipe.tasks.schemaUtils import column_dtype from lsst.utils.timer import timeMethod # Enforce an error for unsafe column/array value setting in pandas. @@ -43,10 +48,29 @@ class AssociationConfig(pexConfig.Config): maxDistArcSeconds = pexConfig.Field( dtype=float, - doc="Maximum distance in arcseconds to test for a DIASource to be a " - "match to a DIAObject.", + doc="Maximum distance in arcseconds for a DIASource to be matched " + "to a DIAObject. This is the sole association radius: every " + "DIAObject within this distance is a match candidate, ranked by " + "position chi^2 when uncertainties are available and by angular " + "distance otherwise.", default=1.0, ) + sigmaFloorArcSeconds = pexConfig.Field( + dtype=float, + doc="Floor on the per-axis position uncertainty (arcsec) used when " + "computing chi^2.", + default=0.05, + ) + fallbackSigmaArcSeconds = pexConfig.Field( + dtype=float, + doc="Per-axis position uncertainty (arcsec) substituted for " + "raErr/decErr values that are missing, non-finite, or " + "non-positive when computing chi^2. Should reflect a plausible " + "size for an unmeasured per-axis uncertainty (e.g., one LSSTCam " + "pixel ~= 0.2 arcsec). Has no effect on rows that carry valid " + "uncertainty values.", + default=0.2, + ) class AssociationTask(pipeBase.Task): @@ -64,7 +88,8 @@ class AssociationTask(pipeBase.Task): @timeMethod def run(self, diaSources, - diaObjects): + diaObjects, + schema=None): """Associate the new DiaSources with existing DiaObjects. Parameters @@ -73,6 +98,10 @@ def run(self, New DIASources to be associated with existing DIAObjects. diaObjects : `pandas.DataFrame` Existing diaObjects from the Apdb. + schema : `dict` [`str`, `felis.datamodel.Schema`] or `None`, optional + Dictionary of Schemas from ``sdm_schemas`` containing the table + definition to use. If `None`, dtypes for new columns are guessed + from the input tables. Returns ------- @@ -100,7 +129,7 @@ def run(self, nUpdatedDiaObjects=0, nUnassociatedDiaObjects=0) - matchResult = self.associate_sources(diaObjects, diaSources) + matchResult = self.associate_sources(diaObjects, diaSources, schema) mask = matchResult.diaSources["diaObjectId"] != 0 @@ -139,7 +168,7 @@ def check_dia_source_radec(self, dia_sources): return dia_sources @timeMethod - def associate_sources(self, dia_objects, dia_sources): + def associate_sources(self, dia_objects, dia_sources, schema=None): """Associate the input DIASources with the catalog of DIAObjects. DiaObject DataFrame must be indexed on ``diaObjectId``. @@ -151,6 +180,9 @@ def associate_sources(self, dia_objects, dia_sources): DIASources into. dia_sources : `pandas.DataFrame` DIASources to associate into the DIAObjectCollection. + schema : `dict` [`str`, `felis.datamodel.Schema`] or `None`, optional + Dictionary of Schemas from ``sdm_schemas`` containing the table + definition to use. Returns ------- @@ -167,72 +199,202 @@ def associate_sources(self, dia_objects, dia_sources): scores = self.score( dia_objects, dia_sources, self.config.maxDistArcSeconds * geom.arcseconds) - match_result = self.match(dia_objects, dia_sources, scores) + match_result = self.match(dia_objects, dia_sources, scores, schema) return match_result @timeMethod def score(self, dia_objects, dia_sources, max_dist): - """Compute a quality score for each dia_source/dia_object pair - between this catalog of DIAObjects and the input DIASource catalog. + """Build the candidate (DIASource, DIAObject) match table and + score every pair. - ``max_dist`` sets maximum separation in arcseconds to consider a - dia_source a possible match to a dia_object. If the pair is - beyond this distance no score is computed. + For each DIASource, all DIAObjects within ``max_dist`` are retrieved + from a kd-tree on unit vectors. Each candidate pair is then scored: + + - If both inputs carry usable ``raErr``/``decErr`` columns, the + score is the 2D Gaussian negative log-likelihood of the + position residual (``0.5 * chi^2 + 0.5 * ln(var_ra * var_dec)``), + so the match prefers the most likely object, not merely the + nearest or the lowest-chi^2 one (see `_position_nll`). + - Otherwise, the distance (in radians) is used as the score. + + No candidates are dropped by the score itself: every pair within + ``max_dist`` is retained, and the score is used only to rank them. + + ``raErr`` and ``decErr`` are taken to follow the LSST DPDD + convention: each is the marginal uncertainty of the catalog + coordinate itself in degrees (no cos(dec) factor folded into + ``raErr``). Under that convention the cos(dec) factor cancels + between residual and uncertainty, and chi^2 reduces to + ``dRA^2 / sum(raErr^2) + dDec^2 / sum(decErr^2)``. + + ``max_dist`` is both the candidate pre-filter and the + association radius: every pair within it is retained, and the + score is used only to rank candidates in the downstream match. Parameters ---------- - dia_objects : `pandas.DataFrame` - A contiguous catalog of DIAObjects to score against dia_sources. - dia_sources : `pandas.DataFrame` - A contiguous catalog of dia_sources to "score" based on distance - and (in the future) other metrics. + dia_objects, dia_sources : `pandas.DataFrame` + Must contain ``ra`` and ``dec``; ``raErr`` and ``decErr`` are + used when present. max_dist : `lsst.geom.Angle` - Maximum allowed distance to compute a score for a given DIAObject - DIASource pair. + Hard angular upper bound on candidate pairs. Returns ------- result : `lsst.pipe.base.Struct` - Results struct with components: - - - ``scores``: array of floats of match quality updated DIAObjects - (array-like of `float`). - - ``obj_idxs``: indexes of the matched DIAObjects in the catalog. - (array-like of `int`) - - ``obj_ids``: array of floats of match quality updated DIAObjects - (array-like of `int`). - - Default values for these arrays are - INF, -1, and -1 respectively for unassociated sources. + Flat candidate-pair table: + + - ``src_idx`` : `numpy.ndarray` of `int` + Positional source index for each surviving pair. + - ``obj_idx`` : `numpy.ndarray` of `int` + Positional object index for each surviving pair. + - ``scores`` : `numpy.ndarray` of `float` + Cost of each pair (position negative log-likelihood if + uncertainty-based, chord distance in radians otherwise). + Lower is better; NLL values may be negative. + - ``unmatched_cost`` : `float` + Cost to assign to the synthetic 'no-match' alternative + in the linear-assignment match — set so that any + surviving real candidate is preferred. """ - scores = np.full(len(dia_sources), np.inf, dtype=np.float64) - obj_idxs = np.full(len(dia_sources), -1, dtype=np.int64) - obj_ids = np.full(len(dia_sources), 0, dtype=np.int64) + n_src = len(dia_sources) + n_obj = len(dia_objects) + empty_int = np.empty(0, dtype=np.int64) + empty_float = np.empty(0, dtype=np.float64) + max_dist_rad = max_dist.asRadians() + # Used as the no-match cost in distance mode; always strictly + # above any real candidate (which the kd-tree caps at max_dist_rad). + chord_unmatched_cost = max_dist_rad * 1.01 + 1e-300 - if len(dia_objects) == 0: + if n_obj == 0 or n_src == 0: return pipeBase.Struct( - scores=scores, - obj_idxs=obj_idxs, - obj_ids=obj_ids) + src_idx=empty_int, + obj_idx=empty_int, + scores=empty_float, + unmatched_cost=chord_unmatched_cost) spatial_tree = self._make_spatial_tree(dia_objects) + src_vectors = self._radec_to_xyz(dia_sources) + + candidate_lists = spatial_tree.query_ball_point( + src_vectors, r=max_dist_rad) + counts = np.fromiter( + (len(c) for c in candidate_lists), dtype=np.int64, count=n_src) + n_pairs = int(counts.sum()) + if n_pairs == 0: + return pipeBase.Struct( + src_idx=empty_int, + obj_idx=empty_int, + scores=empty_float, + unmatched_cost=chord_unmatched_cost) + + src_idx = np.repeat(np.arange(n_src, dtype=np.int64), counts) + obj_idx = np.fromiter( + itertools.chain.from_iterable(candidate_lists), + dtype=np.int64, count=n_pairs) + + if (self._has_position_errors(dia_sources) + and self._has_position_errors(dia_objects)): + scores = self._position_nll(dia_sources, dia_objects, src_idx, obj_idx) + # ``max_dist`` is the sole association gate: every candidate + # pair is kept and the NLL only ranks them. Price the + # 'no-match' alternative just above the worst candidate so a + # real match is always preferred. A diaSource will only not be + # associated if there are more diaSources than diaObjects. + unmatchedCostDelta = 1.0 + unmatched_cost = (float(scores.max()) + unmatchedCostDelta) + else: + obj_vectors = self._radec_to_xyz(dia_objects) + diffs = src_vectors[src_idx] - obj_vectors[obj_idx] + scores = np.linalg.norm(diffs, axis=1) + unmatched_cost = chord_unmatched_cost - max_dist_rad = max_dist.asRadians() + return pipeBase.Struct( + src_idx=src_idx, + obj_idx=obj_idx, + scores=scores, + unmatched_cost=unmatched_cost) - vectors = self._radec_to_xyz(dia_sources) + @staticmethod + def _has_position_errors(catalog): + """Return True iff ``catalog`` carries ``raErr`` and ``decErr`` + columns with at least one finite, positive value in each. + """ + if "raErr" not in catalog.columns or "decErr" not in catalog.columns: + return False + raErr = catalog["raErr"].to_numpy() + decErr = catalog["decErr"].to_numpy() + return (bool(np.any(np.isfinite(raErr) & (raErr > 0.0))) + and bool(np.any(np.isfinite(decErr) & (decErr > 0.0)))) + + def _position_nll(self, dia_sources, dia_objects, src_idx, obj_idx): + """Return the position Gaussian negative log-likelihood (NLL) for + paired DIASources/DIAObjects, used as the match cost. + + Under the hypothesis that the DIASource and DIAObject are the same + astrophysical object, each per-axis residual is a zero-mean + Gaussian whose variance is the sum of the two catalogs' squared + per-axis uncertainties. The cost is the negative log-likelihood of + the observed residual, dropping the constant ``ln(2*pi)`` term: + + NLL = 0.5 * (dra^2 / var_ra + ddec^2 / var_dec) + + 0.5 * ln(var_ra * var_dec). + + The first term is half the 2-DOF position chi^2. The second is the + Gaussian normalization, which penalises poorly-localised + candidates: a bare chi^2 cost omits it, and since a larger + variance deflates chi^2 for a fixed separation, a bare chi^2 would + bias the match toward the more poorly-localised object. + + Non-finite or non-positive per-row uncertainties are replaced + with ``self.config.fallbackSigmaArcSeconds``. The combined per-axis + variance is then floored at ``self.config.sigmaFloorArcSeconds`` + to guard against pathologically small reported errors. - scores, obj_idxs = spatial_tree.query( - vectors, - distance_upper_bound=max_dist_rad) - matched_src_idxs = np.argwhere(np.isfinite(scores)) - obj_ids[matched_src_idxs] = dia_objects.index.to_numpy()[ - obj_idxs[matched_src_idxs]] + Parameters + ---------- + dia_sources, dia_objects : `pandas.DataFrame` + Catalogs containing ``ra``, ``dec``, ``raErr``, ``decErr`` + (all in degrees). + src_idx, obj_idx : `numpy.ndarray` of `int` + Paired positional indices; ``src_idx[k]`` is matched against + ``obj_idx[k]``. - return pipeBase.Struct( - scores=scores, - obj_idxs=obj_idxs, - obj_ids=obj_ids) + Returns + ------- + nll : `numpy.ndarray` of `float` + Position negative log-likelihood, one value per pair. Lower is + a better (more likely) match; values may be negative. + """ + sigma_floor_sq_deg = (self.config.sigmaFloorArcSeconds / 3600.0) ** 2 + fallback_sq_deg = (self.config.fallbackSigmaArcSeconds / 3600.0) ** 2 + + def err_sq(catalog, col, idx): + arr = catalog[col].to_numpy()[idx] + sq = np.square(arr) + return np.where(np.isfinite(sq) & (arr > 0.0), + sq, fallback_sq_deg) + + src_ra = dia_sources["ra"].to_numpy()[src_idx] + src_dec = dia_sources["dec"].to_numpy()[src_idx] + obj_ra = dia_objects["ra"].to_numpy()[obj_idx] + obj_dec = dia_objects["dec"].to_numpy()[obj_idx] + + # Make sure that the RA difference is never greater than +/-180 in + # either direction. + dra = ((src_ra - obj_ra) + 180.0) % 360.0 - 180.0 + ddec = src_dec - obj_dec + + var_ra = (err_sq(dia_sources, "raErr", src_idx) + + err_sq(dia_objects, "raErr", obj_idx)) + var_dec = (err_sq(dia_sources, "decErr", src_idx) + + err_sq(dia_objects, "decErr", obj_idx)) + var_ra = np.maximum(var_ra, sigma_floor_sq_deg) + var_dec = np.maximum(var_dec, sigma_floor_sq_deg) + + chi2 = dra**2/var_ra + ddec**2/var_dec + return 0.5*chi2 + 0.5*np.log(var_ra*var_dec) def _make_spatial_tree(self, dia_objects): """Create a searchable kd-tree the input dia_object positions. @@ -275,72 +437,105 @@ def _radec_to_xyz(self, catalog): return vectors @timeMethod - def match(self, dia_objects, dia_sources, score_struct): - """Match DIAsources to DiaObjects given a score. + def match(self, dia_objects, dia_sources, score_struct, schema=None): + """Solve a min-cost bipartite matching between sources and objects. + + Each DIASource is given a synthetic 'no-match' alternative (a + per-source 'ghost' column) carrying ``score_struct.unmatched_cost``. + A min-weight full bipartite matching is then solved on the sparse + cost matrix made from real candidate pairs and ghost edges. + Sources matched to a ghost are reported as unassociated; sources + matched to a real object inherit that object's ``diaObjectId``. + + When two sources compete for the same object, the source with a + strictly worse alternative gets that object, and the other source falls + back to its second-best candidate rather than creating a new DIAObject. Parameters ---------- - dia_objects : `pandas.DataFrame` - A SourceCatalog of DIAObjects to associate to DIASources. - dia_sources : `pandas.DataFrame` - A contiguous catalog of dia_sources for which the set of scores - has been computed on with DIAObjectCollection.score. + dia_objects, dia_sources : `pandas.DataFrame` + Must contain ``ra`` and ``dec``; ``raErr`` and ``decErr`` are + used when present. score_struct : `lsst.pipe.base.Struct` - Results struct with components: - - - ``"scores"``: array of floats of match quality - updated DIAObjects (array-like of `float`). - - ``"obj_ids"``: array of floats of match quality - updated DIAObjects (array-like of `int`). - - ``"obj_idxs"``: indexes of the matched DIAObjects in the catalog. - (array-like of `int`) - - Default values for these arrays are - INF, -1 and -1 respectively for unassociated sources. + Output of `score`: ``src_idx``, ``obj_idx``, ``scores``, + ``unmatched_cost``. + schema : `dict` [`str`, `felis.datamodel.Schema`] or `None`, optional + Dictionary of Schemas from ``sdm_schemas`` containing the table + definition to use. Returns ------- result : `lsst.pipe.base.Struct` - Results struct with components. - - ``"diaSources"`` : Full set of diaSources both matched and not. - (`pandas.DataFrame`) - - ``"nUpdatedDiaObjects"`` : Number of DiaObjects that were - associated. (`int`) - - ``"nUnassociatedDiaObjects"`` : Number of DiaObjects that were - not matched a new DiaSource. (`int`) + - ``diaSources`` : input source table with ``diaObjectId`` + populated (0 for unmatched). (`pandas.DataFrame`) + - ``nUpdatedDiaObjects`` : number of DIAObjects matched to a + new DIASource. (`int`) + - ``nUnassociatedDiaObjects`` : number of preloaded DIAObjects + with no matching DIASource. (`int`) """ - n_previous_dia_objects = len(dia_objects) - used_dia_object = np.zeros(n_previous_dia_objects, dtype=bool) - used_dia_source = np.zeros(len(dia_sources), dtype=bool) - associated_dia_object_ids = np.zeros(len(dia_sources), - dtype=np.uint64) - n_updated_dia_objects = 0 - - # We sort from best match to worst to effectively perform a - # "handshake" match where both the DIASources and DIAObjects agree - # their the best match. Sources with non-finite scores have no - # match and are skipped — they are left for the new-DIAObject loop. - finite_idx = np.flatnonzero(np.isfinite(score_struct.scores)) - score_args = finite_idx[np.argsort(score_struct.scores[finite_idx])] - for score_idx in score_args: - dia_obj_idx = score_struct.obj_idxs[score_idx] - if used_dia_object[dia_obj_idx]: - continue - used_dia_object[dia_obj_idx] = True - used_dia_source[score_idx] = True - obj_id = score_struct.obj_ids[score_idx] - associated_dia_object_ids[score_idx] = obj_id - n_updated_dia_objects += 1 - - # Assign positionally rather than by label — score_idx values from - # argsort are 0..N-1, but dia_sources may have a non-contiguous - # label index after upstream NaN filtering. + n_src = len(dia_sources) + n_obj = len(dia_objects) + if not pd.api.types.is_integer_dtype(dia_sources["diaObjectId"]): + raise ValueError(f"diaSource column diaObjectId must be an integer, " + f"got {dia_sources['diaObjectId'].dtype} instead") + # Set the diaObjectId dtype from the schema if available. Otherwise + # fall back on the dtype of the incoming diaObjectId column. + # If the dtype is incompatible with the final schema (uint vs int), + # then pandas will silently promote the column to a float. + if schema is not None and schema.get("DiaSource") is not None: + obj_id_col = next( + c for c in schema["DiaSource"].columns if c.name == "diaObjectId" + ) + obj_id_dtype = column_dtype(obj_id_col.datatype, nullable=obj_id_col.nullable) + else: + obj_id_dtype = dia_sources["diaObjectId"].dtype + # Allocate via pandas so pandas-extension dtypes (e.g., "Int64") + # are supported alongside numpy dtypes. + associated_dia_object_ids = pd.array([0]*n_src, dtype=obj_id_dtype) + n_matched = 0 + + if n_src > 0: + # Per-source ghost edges to columns [n_obj, n_obj + n_src). + ghost_src = np.arange(n_src, dtype=np.int64) + ghost_obj = n_obj + ghost_src + ghost_cost = np.full( + n_src, score_struct.unmatched_cost, dtype=np.float64) + + real_weights = score_struct.scores.astype(np.float64, copy=False) + rows = np.concatenate( + [score_struct.src_idx.astype(np.int64, copy=False), + ghost_src]) + cols = np.concatenate( + [score_struct.obj_idx.astype(np.int64, copy=False), + ghost_obj]) + weights = np.concatenate([real_weights, ghost_cost]) + + # scipy's min_weight_full_bipartite_matching treats an explicit + # zero as a missing edge, and the NLL cost can be negative or + # zero. A full matching uses exactly one edge per source, so + # shifting every stored weight by a constant leaves the optimal + # matching unchanged; rescale so the smallest weight is strictly + # positive and no real edge is dropped. + weights = weights - weights.min() + 1.0 + + biadj = csr_matrix( + (weights, (rows, cols)), + shape=(n_src, n_obj + n_src)) + match_src, match_dst = min_weight_full_bipartite_matching(biadj) + + is_real = match_dst < n_obj + matched_src = match_src[is_real] + matched_obj = match_dst[is_real] + n_matched = int(matched_src.size) + if n_matched > 0: + associated_dia_object_ids[matched_src] = ( + dia_objects.index.to_numpy()[matched_obj]) + dia_sources = dia_sources.copy() dia_sources["diaObjectId"] = associated_dia_object_ids return pipeBase.Struct( diaSources=dia_sources, - nUpdatedDiaObjects=n_updated_dia_objects, - nUnassociatedDiaObjects=(n_previous_dia_objects - - n_updated_dia_objects)) + nUpdatedDiaObjects=n_matched, + nUnassociatedDiaObjects=int(n_obj - n_matched)) diff --git a/python/lsst/ap/association/diaPipe.py b/python/lsst/ap/association/diaPipe.py index 5d7c35f0..aba7fa56 100644 --- a/python/lsst/ap/association/diaPipe.py +++ b/python/lsst/ap/association/diaPipe.py @@ -1125,7 +1125,7 @@ def associateDiaSources(self, diaSourceTable, solarSystemObjectTable, diffIm, di associatedSsSources = None unassociatedSsObjects = None # Associate new DiaSources with existing DiaObjects. - assocResults = self.associator.run(unAssocSSDiaSources, diaObjects) + assocResults = self.associator.run(unAssocSSDiaSources, diaObjects, schema=self.schema) createResults = self.createNewDiaObjects(assocResults.unAssocDiaSources) if not assocResults.matchedDiaSources.empty: diff --git a/tests/test_association_task.py b/tests/test_association_task.py index 62873d73..6a62cc86 100644 --- a/tests/test_association_task.py +++ b/tests/test_association_task.py @@ -26,6 +26,7 @@ import lsst.geom as geom import lsst.utils.tests from lsst.ap.association import AssociationTask +from lsst.ap.association.association import AssociationConfig class TestAssociationTask(unittest.TestCase): @@ -106,17 +107,17 @@ def test_score_and_match(self): score_struct = assoc_task.score(self.diaObjects, self.diaSourceZeroScatter, 1.0 * geom.arcseconds) - self.assertFalse(np.isfinite(score_struct.scores[0])) + # Source 0 has no nearby DIAObject (closest is ~144" away). + self.assertNotIn(0, score_struct.src_idx) + # Sources 1..4 each have at least one candidate at ~zero score. for src_idx in range(1, len(self.diaSources)): - # Our scores should be extremely close to 0 but not exactly so due - # to machine noise. - self.assertAlmostEqual(score_struct.scores[src_idx], 0.0, - places=16) - - # After matching each DIAObject should now contain 2 DIASources - # except the last DIAObject in this collection which should be - # newly created during the matching step and contain only one - # DIASource. + mask = score_struct.src_idx == src_idx + self.assertTrue(np.any(mask)) + self.assertAlmostEqual(score_struct.scores[mask].min(), 0.0, + places=10) + + # Linear-assignment match: 4 sources match 4 distinct objects; + # the fifth object is left unassociated. match_result = assoc_task.match( self.diaObjects, self.diaSources, score_struct) self.assertEqual(match_result.nUpdatedDiaObjects, 4) @@ -133,6 +134,285 @@ def test_remove_nan_dia_sources(self): out_dia_sources = assoc_task.check_dia_source_radec(self.diaSources) self.assertEqual(len(out_dia_sources), len(self.diaSources) - 3) + def test_score_falls_back_without_error_columns(self): + """Score uses chord distance when raErr/decErr columns are absent. + """ + task = AssociationTask() + result = task.score(self.diaObjects, self.diaSources, + 1.0 * geom.arcseconds) + # 4 of 5 sources have at least one candidate. + self.assertEqual(len(np.unique(result.src_idx)), 4) + # Chord distance on the unit sphere — values are radians. + self.assertTrue(np.all(result.scores < 1e-3)) + + def test_score_falls_back_when_errors_all_nan(self): + """Empty (all-NaN) raErr/decErr columns trigger the chord fallback. + """ + objects = self.diaObjects.copy() + objects["raErr"] = np.nan + objects["decErr"] = np.nan + sources = self.diaSources.copy() + sources["raErr"] = np.nan + sources["decErr"] = np.nan + task = AssociationTask() + result = task.score(objects, sources, 1.0 * geom.arcseconds) + self.assertEqual(len(np.unique(result.src_idx)), 4) + self.assertTrue(np.all(result.scores < 1e-3)) + + def test_chi2_accepts_within_uncertainty(self): + """A 0.05" separation with 0.05" per-axis sigma yields chi^2 ~ 0.5. + """ + sig_deg = 0.05 / 3600.0 + objects = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0, "raErr": sig_deg, "decErr": sig_deg, + "diaObjectId": 1} + ]).set_index("diaObjectId", drop=False) + sources = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0 + 0.05 / 3600.0, + "raErr": sig_deg, "decErr": sig_deg, + "diaSourceId": 100, "diaObjectId": 0} + ]) + task = AssociationTask() + result = task.score(objects, sources, 1.0 * geom.arcseconds) + self.assertEqual(len(result.src_idx), 1) + self.assertEqual(result.src_idx[0], 0) + self.assertEqual(result.obj_idx[0], 0) + # var per axis = 2 * sig^2 (above the floor); dra=0, ddec=0.05". + # chi^2 = (0.05)^2 / (2 * 0.05^2) = 0.5, and + # NLL = 0.5*chi^2 + 0.5*ln(var_ra * var_dec). + var = 2 * sig_deg ** 2 + expected_nll = 0.5 * 0.5 + 0.5 * np.log(var * var) + self.assertAlmostEqual(result.scores[0], expected_nll, places=6) + + def test_within_maxdist_matches_despite_high_chi2(self): + """A 0.5" separation with 0.05" per-axis sigma (chi^2 ~ 50, i.e. + ~7 sigma) is still matched: the chi^2 only ranks candidates, and + ``maxDistArcSeconds`` is the sole association gate. + """ + sig_deg = 0.05 / 3600.0 + objects = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0, "raErr": sig_deg, "decErr": sig_deg, + "diaObjectId": 1} + ]).set_index("diaObjectId", drop=False) + sources = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0 + sig_deg*10, # 10 sigma in dec + "raErr": sig_deg, "decErr": sig_deg, + "diaSourceId": 100, "diaObjectId": 0} + ]) + task = AssociationTask() + result = task.score(objects, sources, 1.0 * geom.arcseconds) + # The pair is kept as a candidate despite the large offset + # (chi^2 ~ 50) since there is no significance cut. + self.assertEqual(len(result.src_idx), 1) + self.assertTrue(np.isfinite(result.scores[0])) + # End-to-end the source is associated to the object. + run_result = task.run(sources, objects) + self.assertEqual(run_result.nUpdatedDiaObjects, 1) + matched = run_result.matchedDiaSources.set_index("diaSourceId") + self.assertEqual(int(matched.loc[100, "diaObjectId"]), 1) + + def test_nll_prefers_better_localized_object(self): + """With two candidate objects, the NLL cost associates to the + better-localized (tighter) object even though a bare chi^2 would + pick the looser, farther one, whose larger variance deflates its + chi^2. + """ + src_sig = 0.05 / 3600.0 + objects = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0 + 0.15 / 3600.0, # tight and close + "raErr": 0.05 / 3600.0, "decErr": 0.05 / 3600.0, + "diaObjectId": 1}, + {"ra": 1.0, "dec": 1.0 + 0.30 / 3600.0, # loose and far + "raErr": 0.5 / 3600.0, "decErr": 0.5 / 3600.0, + "diaObjectId": 2}, + ]).set_index("diaObjectId", drop=False) + sources = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0, + "raErr": src_sig, "decErr": src_sig, + "diaSourceId": 100, "diaObjectId": 0} + ]) + # A bare chi^2 would prefer the loose object 2 (smaller chi^2)... + chi2_tight = (0.15 / 3600.0) ** 2 / (2 * (0.05 / 3600.0) ** 2) + chi2_loose = (0.30 / 3600.0) ** 2 / ((0.05 / 3600.0) ** 2 + + (0.5 / 3600.0) ** 2) + self.assertLess(chi2_loose, chi2_tight) + + task = AssociationTask() + score_struct = task.score(objects, sources, 1.0 * geom.arcseconds) + # ...but the NLL cost prefers the tight object 1. obj_idx are + # positional: 0 -> diaObjectId 1 (tight), 1 -> diaObjectId 2. + nll = {int(obj): float(score) for obj, score + in zip(score_struct.obj_idx, score_struct.scores)} + self.assertLess(nll[0], nll[1]) + # End-to-end the source associates to the tighter object 1. + result = task.run(sources, objects) + matched = result.matchedDiaSources.set_index("diaSourceId") + self.assertEqual(int(matched.loc[100, "diaObjectId"]), 1) + + def test_chi2_ra_wraparound(self): + """Pairs straddling RA=0/360 are scored correctly. + """ + sig_deg = 0.05 / 3600.0 + objects = pd.DataFrame([ + {"ra": 359.99999, "dec": 0.0, "raErr": sig_deg, "decErr": sig_deg, + "diaObjectId": 1} + ]).set_index("diaObjectId", drop=False) + sources = pd.DataFrame([ + {"ra": 0.00001, "dec": 0.0, + "raErr": sig_deg, "decErr": sig_deg, + "diaSourceId": 100, "diaObjectId": 0} + ]) + task = AssociationTask() + result = task.score(objects, sources, 1.0 * geom.arcseconds) + # Separation is 0.02" of RA across the wrap — within both cuts. + self.assertEqual(len(result.src_idx), 1) + self.assertEqual(result.src_idx[0], 0) + self.assertEqual(result.obj_idx[0], 0) + + def test_lap_recovers_match_lost_to_greedy(self): + """LAP keeps a source matched when its nearest object is taken + by a better-fitting competitor — the case that single-nearest + greedy used to drop into a new DIAObject. + """ + sig_deg = 0.5 / 3600.0 # 0.5" per axis + objects = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0, + "raErr": sig_deg, "decErr": sig_deg, "diaObjectId": 1}, + {"ra": 1.0, "dec": 1.0 + 0.99 / 3600.0, + "raErr": sig_deg, "decErr": sig_deg, "diaObjectId": 2}, + ]).set_index("diaObjectId", drop=False) + # Source 100: closer to obj 1, but obj 2 is also a viable candidate. + # Source 101: sits exactly on obj 1. + sources = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0 + 0.4 / 3600.0, + "raErr": sig_deg, "decErr": sig_deg, "diaSourceId": 100, + "diaObjectId": 0}, + {"ra": 1.0, "dec": 1.0, + "raErr": sig_deg, "decErr": sig_deg, "diaSourceId": 101, + "diaObjectId": 0}, + ]) + task = AssociationTask() + result = task.run(sources, objects) + # Both sources matched; LAP picks src 100 → obj 2, src 101 → obj 1. + self.assertEqual(result.nUpdatedDiaObjects, 2) + self.assertEqual(result.nUnassociatedDiaObjects, 0) + matched = result.matchedDiaSources.set_index("diaSourceId") + self.assertEqual(int(matched.loc[100, "diaObjectId"]), 2) + self.assertEqual(int(matched.loc[101, "diaObjectId"]), 1) + + def test_match_preserves_diaObjectId_dtype(self): + """match() preserves the incoming diaObjectId dtype rather than + overwriting it with uint64. This keeps pd.concat dtype-stable + downstream — mixing uint64 and int64 silently promotes to + float64. + """ + self.assertEqual(self.diaSources["diaObjectId"].dtype, np.int64) + task = AssociationTask() + result = task.run(self.diaSources, self.diaObjects) + self.assertEqual( + result.matchedDiaSources["diaObjectId"].dtype, np.int64) + self.assertEqual( + result.unAssocDiaSources["diaObjectId"].dtype, np.int64) + + sources_u64 = self.diaSources.copy() + sources_u64["diaObjectId"] = sources_u64["diaObjectId"].astype( + np.uint64) + result = task.run(sources_u64, self.diaObjects) + self.assertEqual( + result.matchedDiaSources["diaObjectId"].dtype, np.uint64) + self.assertEqual( + result.unAssocDiaSources["diaObjectId"].dtype, np.uint64) + with self.assertRaises(ValueError): + sources_float = self.diaSources.copy() + sources_float["diaObjectId"] = sources_float["diaObjectId"].astype(float) + task.run(sources_float, self.diaObjects) + + def test_fallback_sigma_scores_missing_error_pair(self): + """A pair whose errors are missing is scored using + ``fallbackSigmaArcSeconds``: the substituted per-axis uncertainty + sets the pair's score. The fallback affects only the score + (ranking), not whether the pair is a candidate. + """ + sig_deg = 0.05 / 3600.0 + # One row in each catalog carries a real error so chi^2 mode + # triggers; the rows under test (id 100 / obj 1) have NaN errors. + objects = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0, + "raErr": np.nan, "decErr": np.nan, "diaObjectId": 1}, + {"ra": 2.0, "dec": 2.0, + "raErr": sig_deg, "decErr": sig_deg, "diaObjectId": 2}, + ]).set_index("diaObjectId", drop=False) + sources = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0 + 0.5 / 3600.0, + "raErr": np.nan, "decErr": np.nan, + "diaSourceId": 100, "diaObjectId": 0}, + {"ra": 2.0, "dec": 2.0, + "raErr": sig_deg, "decErr": sig_deg, + "diaSourceId": 200, "diaObjectId": 0}, + ]) + + def missing_pair_score(fallback): + cfg = AssociationConfig() + cfg.fallbackSigmaArcSeconds = fallback + result = AssociationTask(config=cfg).score( + objects, sources, 1.0 * geom.arcseconds) + # The missing-error pair is source row 0 -> object row 0. + mask = (result.src_idx == 0) & (result.obj_idx == 0) + self.assertEqual(int(mask.sum()), 1) + return float(result.scores[mask][0]) + + def expected_nll(fallback, chi2): + # Both rows have missing errors, so the combined per-axis + # variance is 2 * fallback^2 (deg^2), above the 0.05" floor. + var = 2 * (fallback / 3600.0) ** 2 + return 0.5 * chi2 + 0.5 * np.log(var * var) + + # ddec = 0.5"; fallback 0.05" -> chi^2 = 50, fallback 0.2" -> 3.125. + self.assertAlmostEqual( + missing_pair_score(0.05), expected_nll(0.05, 50.0), places=6) + self.assertAlmostEqual( + missing_pair_score(0.2), expected_nll(0.2, 3.125), places=6) + + # Even with the small fallback (large chi^2) the pair is within + # maxDist, so it is still matched to its object. + cfg = AssociationConfig() + cfg.fallbackSigmaArcSeconds = 0.05 + result = AssociationTask(config=cfg).run(sources, objects) + matched = result.matchedDiaSources.set_index("diaSourceId") + self.assertEqual(int(matched.loc[100, "diaObjectId"]), 1) + self.assertEqual(int(matched.loc[200, "diaObjectId"]), 2) + + def test_contention_leaves_extra_source_unassociated(self): + """When two sources fall within ``maxDistArcSeconds`` of a single + object, the better-fitting one is matched and the other is left + unassociated rather than forced onto the same object. + """ + sig_deg = 0.1 / 3600.0 + objects = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0, + "raErr": sig_deg, "decErr": sig_deg, "diaObjectId": 1}, + ]).set_index("diaObjectId", drop=False) + sources = pd.DataFrame([ + {"ra": 1.0, "dec": 1.0 + 0.2 / 3600.0, # closer to the object + "raErr": sig_deg, "decErr": sig_deg, "diaSourceId": 100, + "diaObjectId": 0}, + {"ra": 1.0, "dec": 1.0 + 0.5 / 3600.0, # farther + "raErr": sig_deg, "decErr": sig_deg, "diaSourceId": 101, + "diaObjectId": 0}, + ]) + task = AssociationTask() + result = task.run(sources, objects) + # One object, so at most one source can match it. + self.assertEqual(result.nUpdatedDiaObjects, 1) + self.assertEqual(result.nUnassociatedDiaObjects, 0) + self.assertEqual(len(result.matchedDiaSources), 1) + self.assertEqual(len(result.unAssocDiaSources), 1) + # The closer source wins the object; the other is unassociated. + matched = result.matchedDiaSources.set_index("diaSourceId") + self.assertEqual(int(matched.loc[100, "diaObjectId"]), 1) + self.assertEqual( + int(result.unAssocDiaSources["diaSourceId"].iloc[0]), 101) + class MemoryTester(lsst.utils.tests.MemoryTestCase): pass diff --git a/tests/test_diaPipe.py b/tests/test_diaPipe.py index 6f1d9f17..2c4644aa 100644 --- a/tests/test_diaPipe.py +++ b/tests/test_diaPipe.py @@ -197,7 +197,7 @@ def solarSystemAssociator_run(unAssocDiaSources, solarSystemObjectTable, visitIn associatedSsSources=_makeMockTable(), unassociatedSsObjects=_makeMockTable()) - def associator_run(table, diaObjects): + def associator_run(table, diaObjects, schema=None): return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3, matchedDiaSources=_makeMockDataFrame(), unAssocDiaSources=_makeMockDataFrame()) diff --git a/tests/utils_tests.py b/tests/utils_tests.py index c3158d54..f960ec9f 100644 --- a/tests/utils_tests.py +++ b/tests/utils_tests.py @@ -103,6 +103,7 @@ def makeDiaSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False bbox = lsst.geom.Box2D(exposure.getBBox()) rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nSources) rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nSources) + diaObjectIds = diaObjectIds.astype(np.int64) if randomizeObjects: objectIds = diaObjectIds[rng.integers(len(diaObjectIds), size=nSources)] else: