From 35460a411264921025ee79d6e6465aec27b639fc Mon Sep 17 00:00:00 2001
From: Ian Sullivan <sullii@uw.edu>
Date: Wed, 6 May 2026 16:36:01 -0700
Subject: [PATCH 1/5] Score DIASource/DIAObject pairs with position chi^2 when
 uncertainties are present

The associator previously ranked candidate matches by chord distance only
and accepted any pair within a fixed 1" cap, ignoring per-source and
per-object positional uncertainties entirely.
When the uncertainty columns are missing or contain no usable values, the
score reverts to the kd-tree chord distance (the previous behavior).
---
 python/lsst/ap/association/association.py | 164 ++++++++++++++++++----
 tests/test_association_task.py            |  90 ++++++++++++
 2 files changed, 223 insertions(+), 31 deletions(-)

diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py
index 0b5b0b13..5ddcc1e8 100644
--- a/python/lsst/ap/association/association.py
+++ b/python/lsst/ap/association/association.py
@@ -43,10 +43,19 @@ 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,
+    )
 
 
 class AssociationTask(pipeBase.Task):
@@ -173,42 +182,57 @@ def associate_sources(self, dia_objects, dia_sources):
 
     @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.
-
-        ``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.
+        """Compute a quality score for each dia_source/dia_object pair.
+
+        The single nearest DIAObject within ``max_dist`` is selected per
+        DIASource via a 3D kd-tree on unit vectors. When both inputs
+        carry usable ``raErr``/``decErr`` columns, the score is the 2-DOF
+        position chi^2 (with non-finite or non-positive uncertainties
+        replaced by ``sigmaFloorArcSeconds`` and the combined per-axis
+        variance also floored at the same value), which is used only to
+        rank the match. ``max_dist`` is the sole association gate. When
+        the uncertainty columns are absent or empty, the kd-tree chord
+        distance (in radians) is used as the score instead — preserving
+        the pre-existing behaviour.
+
+        ``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)``.
 
         Parameters
         ----------
         dia_objects : `pandas.DataFrame`
-            A contiguous catalog of DIAObjects to score against dia_sources.
+            DIAObjects to score against ``dia_sources``. Must contain
+            ``ra`` and ``dec``; ``raErr`` and ``decErr`` are used when
+            present.
         dia_sources : `pandas.DataFrame`
-            A contiguous catalog of dia_sources to "score" based on distance
-            and (in the future) other metrics.
+            DIASources to score. 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.
+            Maximum allowed angular separation; pairs farther than this
+            are not considered regardless of their uncertainties.
 
         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.
+            - ``scores`` : `numpy.ndarray` of `float`
+                Match quality (chi^2 if uncertainty-based, chord
+                distance otherwise). INF for unmatched sources.
+            - ``obj_idxs`` : `numpy.ndarray` of `int`
+                Positional indices of matched DIAObjects; -1 if
+                unmatched.
+            - ``obj_ids`` : `numpy.ndarray` of `int`
+                ``diaObjectId`` of matched DIAObjects; 0 if unmatched.
         """
-        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)
+        scores = np.full(n_src, np.inf, dtype=np.float64)
+        obj_idxs = np.full(n_src, -1, dtype=np.int64)
+        obj_ids = np.full(n_src, 0, dtype=np.int64)
 
         if len(dia_objects) == 0:
             return pipeBase.Struct(
@@ -217,23 +241,101 @@ def score(self, dia_objects, dia_sources, max_dist):
                 obj_ids=obj_ids)
 
         spatial_tree = self._make_spatial_tree(dia_objects)
-
         max_dist_rad = max_dist.asRadians()
-
         vectors = self._radec_to_xyz(dia_sources)
 
-        scores, obj_idxs = spatial_tree.query(
+        chord_dists, candidate_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]]
+        matched = np.isfinite(chord_dists)
+        if not np.any(matched):
+            return pipeBase.Struct(
+                scores=scores,
+                obj_idxs=obj_idxs,
+                obj_ids=obj_ids)
+
+        use_chi2 = (self._has_position_errors(dia_sources)
+                    and self._has_position_errors(dia_objects))
+        if use_chi2:
+            src_idx = np.flatnonzero(matched)
+            obj_idx = candidate_obj_idxs[src_idx]
+            chi2 = self._chi2_position(
+                dia_sources, dia_objects, src_idx, obj_idx)
+            scores[src_idx] = chi2
+            obj_idxs[src_idx] = obj_idx
+            obj_ids[src_idx] = dia_objects.index.to_numpy()[obj_idx]
+        else:
+            scores[matched] = chord_dists[matched]
+            obj_idxs[matched] = candidate_obj_idxs[matched]
+            obj_ids[matched] = dia_objects.index.to_numpy()[
+                candidate_obj_idxs[matched]]
 
         return pipeBase.Struct(
             scores=scores,
             obj_idxs=obj_idxs,
             obj_ids=obj_ids)
 
+    @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 _chi2_position(self, dia_sources, dia_objects, src_idx, obj_idx):
+        """Return the 2-DOF position chi^2 for paired DIASources/DIAObjects.
+
+        Non-finite or non-positive per-row uncertainties are replaced
+        with ``self.config.sigmaFloorArcSeconds``; the combined per-axis
+        variance is also floored at that same value to guard against
+        pathologically small reported errors.
+
+        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]``.
+
+        Returns
+        -------
+        chi2 : `numpy.ndarray` of `float`
+             2 degrees of freedom position chi^2, one value per pair.
+        """
+        sigma_floor_sq_deg = (self.config.sigmaFloorArcSeconds / 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, sigma_floor_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)
+
+        return dra * dra / var_ra + ddec * ddec / var_dec
+
     def _make_spatial_tree(self, dia_objects):
         """Create a searchable kd-tree the input dia_object positions.
 
diff --git a/tests/test_association_task.py b/tests/test_association_task.py
index 62873d73..adbba7bd 100644
--- a/tests/test_association_task.py
+++ b/tests/test_association_task.py
@@ -133,6 +133,96 @@ 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 match (source 0 has no nearby object).
+        self.assertEqual(np.sum(np.isfinite(result.scores)), 4)
+        finite = result.scores[np.isfinite(result.scores)]
+        # Chord distance on the unit sphere — values are radians.
+        self.assertTrue(np.all(finite < 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)
+        finite = result.scores[np.isfinite(result.scores)]
+        self.assertEqual(len(finite), 4)
+        # Chord-distance regime, not chi^2.
+        self.assertTrue(np.all(finite < 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.assertTrue(np.isfinite(result.scores[0]))
+        # var per axis = 2 * sig^2 = 2*(0.05)^2 arcsec^2; dra=0, ddec=0.05".
+        # chi^2 = 0 + (0.05)^2 / (2 * 0.05^2) = 0.5.
+        self.assertAlmostEqual(result.scores[0], 0.5, places=6)
+
+    def test_within_maxdist_matches_despite_high_chi2(self):
+        """A 0.5" separation with 0.05" per-axis sigma (chi^2 ~ 50) is
+        still matched: the chi^2 only ranks the match, 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)
+        # Within the 1" gate, so it is scored (chi^2 ~ 50) and matched.
+        self.assertTrue(np.isfinite(result.scores[0]))
+        self.assertAlmostEqual(result.scores[0], 50.0, places=6)
+        self.assertEqual(result.obj_ids[0], 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; at dec=0 that's ~0.02"
+        # on-sky, well within both the geometric and chi^2 cuts.
+        self.assertTrue(np.isfinite(result.scores[0]))
+        self.assertEqual(result.obj_ids[0], 1)
+
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):
     pass

From 9e2b4f65e15d2cea416c78cc8cb0223383a1bfbc Mon Sep 17 00:00:00 2001
From: Ian Sullivan <sullii@uw.edu>
Date: Wed, 6 May 2026 16:55:28 -0700
Subject: [PATCH 2/5] Solve DIASource/DIAObject association as a min-cost
 bipartite matching

The previous matcher considered only each DIASource's single nearest
DIAObject. When two sources competed for the same object, the loser
would be unassociated and triggered a new DIAObject even if a viable
second-best candidate was nearby.
---
 python/lsst/ap/association/association.py | 267 ++++++++++++----------
 tests/test_association_task.py            | 149 +++++++++---
 tests/utils_tests.py                      |   1 +
 3 files changed, 273 insertions(+), 144 deletions(-)

diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py
index 5ddcc1e8..724eef83 100644
--- a/python/lsst/ap/association/association.py
+++ b/python/lsst/ap/association/association.py
@@ -24,8 +24,12 @@
 
 __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
@@ -182,18 +186,19 @@ def associate_sources(self, dia_objects, dia_sources):
 
     @timeMethod
     def score(self, dia_objects, dia_sources, max_dist):
-        """Compute a quality score for each dia_source/dia_object pair.
-
-        The single nearest DIAObject within ``max_dist`` is selected per
-        DIASource via a 3D kd-tree on unit vectors. When both inputs
-        carry usable ``raErr``/``decErr`` columns, the score is the 2-DOF
-        position chi^2 (with non-finite or non-positive uncertainties
-        replaced by ``sigmaFloorArcSeconds`` and the combined per-axis
-        variance also floored at the same value), which is used only to
-        rank the match. ``max_dist`` is the sole association gate. When
-        the uncertainty columns are absent or empty, the kd-tree chord
-        distance (in radians) is used as the score instead — preserving
-        the pre-existing behaviour.
+        """Build the candidate (DIASource, DIAObject) match table and
+        score every pair.
+
+        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 2-DOF position chi^2, so the match prefers the
+          best-fitting object rather than the merely-nearest one.
+        - 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
@@ -202,78 +207,93 @@ def score(self, dia_objects, dia_sources, max_dist):
         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`
-            DIAObjects to score against ``dia_sources``. Must contain
-            ``ra`` and ``dec``; ``raErr`` and ``decErr`` are used when
-            present.
-        dia_sources : `pandas.DataFrame`
-            DIASources to score. Must contain ``ra`` and ``dec``;
-            ``raErr`` and ``decErr`` are used when present.
+        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 angular separation; pairs farther than this
-            are not considered regardless of their uncertainties.
+            Hard angular upper bound on candidate pairs.
 
         Returns
         -------
         result : `lsst.pipe.base.Struct`
-            Results struct with components:
+            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`
-                Match quality (chi^2 if uncertainty-based, chord
-                distance otherwise). INF for unmatched sources.
-            - ``obj_idxs`` : `numpy.ndarray` of `int`
-                Positional indices of matched DIAObjects; -1 if
-                unmatched.
-            - ``obj_ids`` : `numpy.ndarray` of `int`
-                ``diaObjectId`` of matched DIAObjects; 0 if unmatched.
+                Cost of each pair (chi^2 if uncertainty-based, chord
+                distance in radians otherwise). Lower is better.
+            - ``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.
         """
         n_src = len(dia_sources)
-        scores = np.full(n_src, np.inf, dtype=np.float64)
-        obj_idxs = np.full(n_src, -1, dtype=np.int64)
-        obj_ids = np.full(n_src, 0, dtype=np.int64)
+        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)
-        max_dist_rad = max_dist.asRadians()
-        vectors = self._radec_to_xyz(dia_sources)
-
-        chord_dists, candidate_obj_idxs = spatial_tree.query(
-            vectors,
-            distance_upper_bound=max_dist_rad)
-        matched = np.isfinite(chord_dists)
-        if not np.any(matched):
+        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(
-                scores=scores,
-                obj_idxs=obj_idxs,
-                obj_ids=obj_ids)
-
-        use_chi2 = (self._has_position_errors(dia_sources)
-                    and self._has_position_errors(dia_objects))
-        if use_chi2:
-            src_idx = np.flatnonzero(matched)
-            obj_idx = candidate_obj_idxs[src_idx]
-            chi2 = self._chi2_position(
+                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._chi2_position(
                 dia_sources, dia_objects, src_idx, obj_idx)
-            scores[src_idx] = chi2
-            obj_idxs[src_idx] = obj_idx
-            obj_ids[src_idx] = dia_objects.index.to_numpy()[obj_idx]
+            # ``max_dist`` is the sole association gate: every candidate
+            # pair is kept and the chi^2 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:
-            scores[matched] = chord_dists[matched]
-            obj_idxs[matched] = candidate_obj_idxs[matched]
-            obj_ids[matched] = dia_objects.index.to_numpy()[
-                candidate_obj_idxs[matched]]
+            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
 
         return pipeBase.Struct(
+            src_idx=src_idx,
+            obj_idx=obj_idx,
             scores=scores,
-            obj_idxs=obj_idxs,
-            obj_ids=obj_ids)
+            unmatched_cost=unmatched_cost)
 
     @staticmethod
     def _has_position_errors(catalog):
@@ -378,71 +398,86 @@ def _radec_to_xyz(self, catalog):
 
     @timeMethod
     def match(self, dia_objects, dia_sources, score_struct):
-        """Match DIAsources to DiaObjects given a score.
+        """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`
         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``.
 
         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")
+        # Preserve the dtype of the incoming diaObjectId column so the
+        # output is concat-stable downstream (mixing uint64 and int64
+        # forces a silent promotion to float64 in pd.concat).
+        associated_dia_object_ids = np.zeros(n_src, dtype=dia_sources["diaObjectId"].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)
+
+            # scipy's min_weight_full_bipartite_matching removes
+            # explicit-zero edges before solving. Nudge them off zero
+            # so that perfect-match candidates (chi^2 = 0) are still
+            # seen by the LAP.
+            real_weights = np.maximum(
+                score_struct.scores.astype(np.float64, copy=False), 1e-300)
+            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])
+
+            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/tests/test_association_task.py b/tests/test_association_task.py
index adbba7bd..61ab863c 100644
--- a/tests/test_association_task.py
+++ b/tests/test_association_task.py
@@ -106,17 +106,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)
@@ -139,11 +139,10 @@ def test_score_falls_back_without_error_columns(self):
         task = AssociationTask()
         result = task.score(self.diaObjects, self.diaSources,
                             1.0 * geom.arcseconds)
-        # 4 of 5 sources match (source 0 has no nearby object).
-        self.assertEqual(np.sum(np.isfinite(result.scores)), 4)
-        finite = result.scores[np.isfinite(result.scores)]
+        # 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(finite < 1e-3))
+        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.
@@ -156,10 +155,8 @@ def test_score_falls_back_when_errors_all_nan(self):
         sources["decErr"] = np.nan
         task = AssociationTask()
         result = task.score(objects, sources, 1.0 * geom.arcseconds)
-        finite = result.scores[np.isfinite(result.scores)]
-        self.assertEqual(len(finite), 4)
-        # Chord-distance regime, not chi^2.
-        self.assertTrue(np.all(finite < 1e-3))
+        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.
@@ -176,14 +173,16 @@ def test_chi2_accepts_within_uncertainty(self):
         ])
         task = AssociationTask()
         result = task.score(objects, sources, 1.0 * geom.arcseconds)
-        self.assertTrue(np.isfinite(result.scores[0]))
+        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 = 2*(0.05)^2 arcsec^2; dra=0, ddec=0.05".
         # chi^2 = 0 + (0.05)^2 / (2 * 0.05^2) = 0.5.
         self.assertAlmostEqual(result.scores[0], 0.5, places=6)
 
     def test_within_maxdist_matches_despite_high_chi2(self):
-        """A 0.5" separation with 0.05" per-axis sigma (chi^2 ~ 50) is
-        still matched: the chi^2 only ranks the match, and
+        """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
@@ -198,10 +197,15 @@ def test_within_maxdist_matches_despite_high_chi2(self):
         ])
         task = AssociationTask()
         result = task.score(objects, sources, 1.0 * geom.arcseconds)
-        # Within the 1" gate, so it is scored (chi^2 ~ 50) and matched.
-        self.assertTrue(np.isfinite(result.scores[0]))
+        # The pair is kept as a candidate and scored at its chi^2 (~50),
+        # despite the large offset.
+        self.assertEqual(len(result.src_idx), 1)
         self.assertAlmostEqual(result.scores[0], 50.0, places=6)
-        self.assertEqual(result.obj_ids[0], 1)
+        # 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_chi2_ra_wraparound(self):
         """Pairs straddling RA=0/360 are scored correctly.
@@ -218,10 +222,99 @@ def test_chi2_ra_wraparound(self):
         ])
         task = AssociationTask()
         result = task.score(objects, sources, 1.0 * geom.arcseconds)
-        # Separation is 0.02" of RA across the wrap; at dec=0 that's ~0.02"
-        # on-sky, well within both the geometric and chi^2 cuts.
-        self.assertTrue(np.isfinite(result.scores[0]))
-        self.assertEqual(result.obj_ids[0], 1)
+        # 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_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):
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:

From cf6b2b351f68f3c08af0e56096c734bbdae3b22e Mon Sep 17 00:00:00 2001
From: Ian Sullivan <sullii@uw.edu>
Date: Mon, 11 May 2026 15:56:06 -0700
Subject: [PATCH 3/5] Use DiaSource schema to set new column dtypes in
 association

---
 python/lsst/ap/association/association.py | 42 ++++++++++++++++++-----
 python/lsst/ap/association/diaPipe.py     |  2 +-
 tests/test_diaPipe.py                     |  2 +-
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py
index 724eef83..a727338f 100644
--- a/python/lsst/ap/association/association.py
+++ b/python/lsst/ap/association/association.py
@@ -35,6 +35,7 @@
 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.
@@ -77,7 +78,8 @@ class AssociationTask(pipeBase.Task):
     @timeMethod
     def run(self,
             diaSources,
-            diaObjects):
+            diaObjects,
+            schema=None):
         """Associate the new DiaSources with existing DiaObjects.
 
         Parameters
@@ -86,6 +88,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
         -------
@@ -113,7 +119,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
 
@@ -152,7 +158,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``.
@@ -164,6 +170,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
         -------
@@ -180,7 +189,7 @@ 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
 
@@ -397,7 +406,7 @@ def _radec_to_xyz(self, catalog):
         return vectors
 
     @timeMethod
-    def match(self, dia_objects, dia_sources, score_struct):
+    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
@@ -414,9 +423,14 @@ def match(self, dia_objects, dia_sources, score_struct):
         Parameters
         ----------
         dia_objects, dia_sources : `pandas.DataFrame`
+            Must contain ``ra`` and ``dec``; ``raErr`` and ``decErr`` are
+            used when present.
         score_struct : `lsst.pipe.base.Struct`
             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
         -------
@@ -434,10 +448,20 @@ def match(self, dia_objects, dia_sources, score_struct):
         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")
-        # Preserve the dtype of the incoming diaObjectId column so the
-        # output is concat-stable downstream (mixing uint64 and int64
-        # forces a silent promotion to float64 in pd.concat).
-        associated_dia_object_ids = np.zeros(n_src, dtype=dia_sources["diaObjectId"].dtype)
+        # 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:
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_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())

From 5c0a1a1f5d0fdb8b1c053d3cf9c92372c03f37ef Mon Sep 17 00:00:00 2001
From: Ian Sullivan <sullii@uw.edu>
Date: Thu, 21 May 2026 13:56:47 -0700
Subject: [PATCH 4/5] Add falback position uncertainty for sources where it is
 missing

---
 python/lsst/ap/association/association.py | 19 +++++++--
 tests/test_association_task.py            | 49 +++++++++++++++++++++++
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py
index a727338f..90b0863e 100644
--- a/python/lsst/ap/association/association.py
+++ b/python/lsst/ap/association/association.py
@@ -61,6 +61,16 @@ class AssociationConfig(pexConfig.Config):
         "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):
@@ -320,9 +330,9 @@ def _chi2_position(self, dia_sources, dia_objects, src_idx, obj_idx):
         """Return the 2-DOF position chi^2 for paired DIASources/DIAObjects.
 
         Non-finite or non-positive per-row uncertainties are replaced
-        with ``self.config.sigmaFloorArcSeconds``; the combined per-axis
-        variance is also floored at that same value to guard against
-        pathologically small reported errors.
+        with ``self.config.fallbackSigmaArcSeconds``. The combined per-axis
+        variance is then floored at ``self.config.sigmaFloorArcSeconds``
+        to guard against pathologically small reported errors.
 
         Parameters
         ----------
@@ -339,12 +349,13 @@ def _chi2_position(self, dia_sources, dia_objects, src_idx, obj_idx):
              2 degrees of freedom position chi^2, one value per pair.
         """
         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, sigma_floor_sq_deg)
+                            sq, fallback_sq_deg)
 
         src_ra = dia_sources["ra"].to_numpy()[src_idx]
         src_dec = dia_sources["dec"].to_numpy()[src_idx]
diff --git a/tests/test_association_task.py b/tests/test_association_task.py
index 61ab863c..da13cacf 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):
@@ -285,6 +286,54 @@ def test_match_preserves_diaObjectId_dtype(self):
             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``: a larger fallback yields a smaller
+        chi^2 for the same separation.
+        """
+        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])
+
+        # var_dec = 2 * fallback^2 (floored at sigmaFloor=0.05"); ddec=0.5".
+        # fallback 0.05" -> chi^2 = 0.25 / (2*0.05^2) = 50.
+        self.assertAlmostEqual(missing_pair_score(0.05), 50.0, places=6)
+        # fallback 0.2"  -> chi^2 = 0.25 / (2*0.2^2)  = 3.125.
+        self.assertAlmostEqual(missing_pair_score(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

From 5cf88d8e38c776fddbf9e9f79135696c24b0bef6 Mon Sep 17 00:00:00 2001
From: Ian Sullivan <sullii@uw.edu>
Date: Tue, 23 Jun 2026 14:11:45 -0700
Subject: [PATCH 5/5] Use a negative log-liklihood function to score competing
 diaObject associations

This avoids the case where a further diaObject with worse coordinate errors has a lower chi-squared than a nearby diaObject with tight errors.
---
 python/lsst/ap/association/association.py | 59 ++++++++++++------
 tests/test_association_task.py            | 74 +++++++++++++++++++----
 2 files changed, 102 insertions(+), 31 deletions(-)

diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py
index 90b0863e..4ca0329e 100644
--- a/python/lsst/ap/association/association.py
+++ b/python/lsst/ap/association/association.py
@@ -212,8 +212,10 @@ def score(self, dia_objects, dia_sources, max_dist):
         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 2-DOF position chi^2, so the match prefers the
-          best-fitting object rather than the merely-nearest one.
+          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
@@ -248,8 +250,9 @@ def score(self, dia_objects, dia_sources, max_dist):
             - ``obj_idx`` : `numpy.ndarray` of `int`
                 Positional object index for each surviving pair.
             - ``scores`` : `numpy.ndarray` of `float`
-                Cost of each pair (chi^2 if uncertainty-based, chord
-                distance in radians otherwise). Lower is better.
+                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
@@ -293,10 +296,9 @@ def score(self, dia_objects, dia_sources, max_dist):
 
         if (self._has_position_errors(dia_sources)
                 and self._has_position_errors(dia_objects)):
-            scores = self._chi2_position(
-                dia_sources, dia_objects, src_idx, obj_idx)
+            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 chi^2 only ranks them. Price the
+            # 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.
@@ -326,8 +328,24 @@ def _has_position_errors(catalog):
         return (bool(np.any(np.isfinite(raErr) & (raErr > 0.0)))
                 and bool(np.any(np.isfinite(decErr) & (decErr > 0.0))))
 
-    def _chi2_position(self, dia_sources, dia_objects, src_idx, obj_idx):
-        """Return the 2-DOF position chi^2 for paired DIASources/DIAObjects.
+    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
@@ -345,8 +363,9 @@ def _chi2_position(self, dia_sources, dia_objects, src_idx, obj_idx):
 
         Returns
         -------
-        chi2 : `numpy.ndarray` of `float`
-             2 degrees of freedom position chi^2, one value per pair.
+        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
@@ -374,7 +393,8 @@ def err_sq(catalog, col, idx):
         var_ra = np.maximum(var_ra, sigma_floor_sq_deg)
         var_dec = np.maximum(var_dec, sigma_floor_sq_deg)
 
-        return dra * dra / var_ra + ddec * ddec / var_dec
+        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.
@@ -482,12 +502,7 @@ def match(self, dia_objects, dia_sources, score_struct, schema=None):
             ghost_cost = np.full(
                 n_src, score_struct.unmatched_cost, dtype=np.float64)
 
-            # scipy's min_weight_full_bipartite_matching removes
-            # explicit-zero edges before solving. Nudge them off zero
-            # so that perfect-match candidates (chi^2 = 0) are still
-            # seen by the LAP.
-            real_weights = np.maximum(
-                score_struct.scores.astype(np.float64, copy=False), 1e-300)
+            real_weights = score_struct.scores.astype(np.float64, copy=False)
             rows = np.concatenate(
                 [score_struct.src_idx.astype(np.int64, copy=False),
                  ghost_src])
@@ -496,6 +511,14 @@ def match(self, dia_objects, dia_sources, score_struct, schema=None):
                  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))
diff --git a/tests/test_association_task.py b/tests/test_association_task.py
index da13cacf..6a62cc86 100644
--- a/tests/test_association_task.py
+++ b/tests/test_association_task.py
@@ -177,9 +177,12 @@ def test_chi2_accepts_within_uncertainty(self):
         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 = 2*(0.05)^2 arcsec^2; dra=0, ddec=0.05".
-        # chi^2 = 0 + (0.05)^2 / (2 * 0.05^2) = 0.5.
-        self.assertAlmostEqual(result.scores[0], 0.5, places=6)
+        # 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.
@@ -198,16 +201,54 @@ def test_within_maxdist_matches_despite_high_chi2(self):
         ])
         task = AssociationTask()
         result = task.score(objects, sources, 1.0 * geom.arcseconds)
-        # The pair is kept as a candidate and scored at its chi^2 (~50),
-        # despite the large offset.
+        # 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.assertAlmostEqual(result.scores[0], 50.0, places=6)
+        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.
         """
@@ -288,8 +329,9 @@ def test_match_preserves_diaObjectId_dtype(self):
 
     def test_fallback_sigma_scores_missing_error_pair(self):
         """A pair whose errors are missing is scored using
-        ``fallbackSigmaArcSeconds``: a larger fallback yields a smaller
-        chi^2 for the same separation.
+        ``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
@@ -319,11 +361,17 @@ def missing_pair_score(fallback):
             self.assertEqual(int(mask.sum()), 1)
             return float(result.scores[mask][0])
 
-        # var_dec = 2 * fallback^2 (floored at sigmaFloor=0.05"); ddec=0.5".
-        # fallback 0.05" -> chi^2 = 0.25 / (2*0.05^2) = 50.
-        self.assertAlmostEqual(missing_pair_score(0.05), 50.0, places=6)
-        # fallback 0.2"  -> chi^2 = 0.25 / (2*0.2^2)  = 3.125.
-        self.assertAlmostEqual(missing_pair_score(0.2), 3.125, places=6)
+        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.