From 339e6eaaf460c625978876247d38f5d119d74d38 Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Wed, 21 Jan 2026 15:27:11 -0800 Subject: [PATCH 1/9] Add diaObject deduplication task --- python/lsst/ap/association/__init__.py | 1 + .../deduplicateAllSkyDiaObjects.py | 365 ++++++++++++++++++ tests/test_deduplicateAllSkyDiaObjects.py | 298 ++++++++++++++ 3 files changed, 664 insertions(+) create mode 100644 python/lsst/ap/association/deduplicateAllSkyDiaObjects.py create mode 100644 tests/test_deduplicateAllSkyDiaObjects.py diff --git a/python/lsst/ap/association/__init__.py b/python/lsst/ap/association/__init__.py index 2822a375..3e0958ff 100644 --- a/python/lsst/ap/association/__init__.py +++ b/python/lsst/ap/association/__init__.py @@ -25,6 +25,7 @@ from .loadDiaCatalogs import * from .packageAlerts import * from .diaPipe import * +from .deduplicateAllSkyDiaObjects import * from .exportDiaCatalogs import * from .transformDiaSourceCatalog import * from .utils import * diff --git a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py new file mode 100644 index 00000000..d2db91b6 --- /dev/null +++ b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py @@ -0,0 +1,365 @@ +# This file is part of ap_association. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +from astropy.coordinates import SkyCoord +from astropy.table import Table +from astropy.time import Time +import astropy.units as u +import numpy as np +import pandas as pd +from sklearn.cluster import AgglomerativeClustering +from sklearn.neighbors import kneighbors_graph + +import lsst.dax.apdb as daxApdb +import lsst.pex.config as pexConfig +import lsst.pipe.base as pipeBase + +__all__ = ("DeduplicateAllSkyDiaObjectsTask", "DeduplicateAllSkyDiaObjectsConfig") + + +class DeduplicateAllSkyDiaObjectsConfig(pexConfig.Config): + """Configuration for DeduplicateAllSkyDiaObjectsTask. + + Parameters + ---------- + maxClusteringDistance : `float` + Maximum distance to merge clusters of duplicate DiaObjects, in + arcseconds. Default is 1.5 arcseconds. + earliestMidpointMjdTai : `float` + MidpointMjdTai (Modified Julian Date in TAI scale) of the + earliest DiaSource to be reassigned during deduplication. + Default corresponds to 2025-09-01. + maxDiaObjects : `int` + Maximum number of DiaObjects to process. An exception is raised + if more DiaObjects are found to prevent excessive memory/computation. + """ + apdb_config_url = pexConfig.Field( + dtype=str, + default=None, + optional=False, + doc="A config file specifying the APDB and its connection parameters, " + "typically written by the apdb-cli command-line utility. " + "The database must already be initialized.", + ) + maxClusteringDistance = pexConfig.RangeField( + doc="Maximum distence to merge clusters of duplicate diaObjects (arcseconds)", + dtype=float, + default=1.5, + min=0, + ) + earliestMidpointMjdTai = pexConfig.Field( + dtype=float, + default=Time('2025-09-01', scale='tai').mjd, + doc="MidpointMjdTai of earliest DIASource to be reassigned during " + "deduplication.", + ) + maxDiaObjects = pexConfig.RangeField( + doc="Maximum number of DiaObjects to process. An exception is raised " + "if more DiaObjects are found to prevent excessive memory/computation.", + dtype=int, + default=1000000, + min=1, + ) + nNeighborsConnectivity = pexConfig.RangeField( + doc="Number of neighbors to use for clustering. Larger values are more" + "accurate but more expensive computationally.", + dtype=int, + default=30, + min=3, + ) + + +class DeduplicateAllSkyDiaObjectsTask(pipeBase.Task): + """Deduplicate DiaObjects by clustering spatially and update the APDB. + + This task identifies and merges duplicate DiaObjects in the Alert + Production Database (APDB) that are spatially coincident. It uses + agglomerative clustering to group nearby objects, then reassigns + DiaSources from duplicate objects to a single kept DiaObject (the + oldest one in each cluster). + """ + ConfigClass = DeduplicateAllSkyDiaObjectsConfig + _DefaultName = "deduplicateAllSkyDiaObjects" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.apdb = daxApdb.Apdb.from_uri(self.config.apdb_config_url) + + def run(self): + """Load DiaObjects and cluster them to remove duplicates. + + Returns + ------- + result : `lsst.pipe.base.Struct` + Results struct with components: + + ``diaObjectDeduplicationMap`` : `astropy.table.Table` + Table mapping duplicate DiaObject IDs to the IDs they + should be merged into. Contains columns + 'removedDiaObjectId' and 'keptDiaObjectId'. + + Raises + ------ + lsst.pipe.base.NoWorkFound + If no duplicate DiaObjects are found in the APDB. + """ + + # Return all DiaObject versions created since the last deduplication + # run--typically this will be since yesterday morning. + # TODO: consider adding a config for the "since" argument + all_diaObjects = self.apdb.getDiaObjectsForDedup() + # fix getDiaObjectsForDedup returning a non-unique pandas index + all_diaObjects = all_diaObjects.reset_index().drop('index', axis=1) + + if len(all_diaObjects) > self.config.maxDiaObjects: + raise RuntimeError(f"Found {len(all_diaObjects)} DiaObjects, which exceeds the " + f"configured maximum of {self.config.maxDiaObjects}. " + "Aborting to avoid excessive memory/computation.") + self.log.info(f"Loaded {len(all_diaObjects)} DiaObjects.") + + # only keep the latest versions of the DIAObjects + all_diaObjects.sort_values(['diaObjectId', 'validityStartMjdTai'], inplace=True) + wdup = all_diaObjects.duplicated(subset='diaObjectId', keep='last') + # copy to avoid setting on copy errors + diaObjects = all_diaObjects.loc[~wdup, :].copy() + self.log.info(f"Filtered to {len(all_diaObjects)} latest DiaObjects.") + + duplicate_count = self.count_duplicates(diaObjects) + + if duplicate_count == 0: + raise pipeBase.NoWorkFound("No duplicate DiaObjects found.") + else: + self.log.info(f"Found {duplicate_count} duplicates in {len(diaObjects)} DiaObjects.") + + cluster_labels = self.cluster(diaObjects) + + diaObjects.loc[:, 'cluster_label'] = cluster_labels + + deduplication_map = self.remap_clusters(diaObjects) + + self.log.info(f"Reassigned {len(deduplication_map)} duplicates to " + f"{len(np.unique(deduplication_map['keptDiaObjectId']))} DiaObjects.") + + removed_ids = deduplication_map['removedDiaObjectId'].to_list() + wremoved = diaObjects['diaObjectId'].apply(lambda x: x in removed_ids) + kept_diaObjects = diaObjects.loc[~wremoved, :] + updated_duplicate_count = self.count_duplicates(kept_diaObjects) + self.log.info(f"After deduplication, {updated_duplicate_count} duplicates " + f"remain of {len(kept_diaObjects)} DiaObjects.") + + # TODO: consider a "dry run" configuration + self.remove_apdb_duplicates(diaObjects, deduplication_map) + + return pipeBase.Struct( + diaObjectDeduplicationMap=Table.from_df(deduplication_map)) + + def remove_apdb_duplicates(self, diaObjects, deduplication_map): + """Reassign DiaSources and remove DiaObjects per provided map. + + This method performs the actual database modifications to + deduplicate DiaObjects. It reassigns DiaSources from duplicate + DiaObjects to their corresponding kept objects, then marks the + duplicate DiaObjects as invalid by setting their validity end + time. + + Parameters + ---------- + diaObjects : `pandas.DataFrame` + DataFrame containing DiaObjects, including both kept and + removed objects. Must include columns: 'diaObjectId', + 'validityStartMjdTai', 'ra', 'dec'. + deduplication_map : `pandas.DataFrame` + Table specifying how to reassociate duplicates. Must contain + columns: + - 'removedDiaObjectId' : IDs of DiaObjects to be removed + - 'keptDiaObjectId' : IDs of DiaObjects to keep + """ + + start_time = Time(self.config.earliestMidpointMjdTai, + format='mjd', scale='tai') + + # make the records needed by the apdb methods + DiaObjectsToRemove = [] + for idx, idi in deduplication_map['removedDiaObjectId'].items(): + w = diaObjects['diaObjectId'] == idi + assert (np.sum(w) == 1) + DiaObjectsToRemove.extend([daxApdb.recordIds.DiaObjectId.from_named_tuple(row) + for row in diaObjects.loc[w].itertuples()]) + + diaSourcesToReassign = \ + self.apdb.getDiaSourcesForDiaObjects(DiaObjectsToRemove, + start_time=start_time, + max_dist_arcsec=self.config.maxClusteringDistance) + + id_map = {} + for old_id, new_id in zip(deduplication_map['removedDiaObjectId'], + deduplication_map['keptDiaObjectId']): + w = diaSourcesToReassign['diaObjectId'] == old_id + reassign_count = np.sum(w) + + if reassign_count: + for row in diaSourcesToReassign.loc[w].itertuples(): + ds = daxApdb.recordIds.DiaSourceId.from_named_tuple(row) + id_map[ds] = new_id + self.log.verbose('Reassigned %d diaSources from diaObject %d to %d' % + (reassign_count, old_id, new_id)) + + else: + self.log.verbose('No diaSources found for diaObject %d' % old_id) + + self.apdb.reassignDiaSourcesToDiaObjects(id_map) + + self.apdb.setValidityEnd(DiaObjectsToRemove, Time.now()) + + # Clear the table staging DiaObjects for deduplication + self.apdb.resetDedup() + + def remap_clusters(self, diaObjects): + """Use cluster labels to determine which DiaObjects to remap. + + For each cluster containing multiple DiaObjects, this method + identifies the oldest object (by validityStartMjdTai) to keep + and marks all others for removal. + + Parameters + ---------- + diaObjects : `pandas.DataFrame` + DataFrame of DiaObjects with cluster assignments. Must + include columns: + + - 'cluster_label' : Cluster ID from clustering algorithm + - 'diaObjectId' : Unique identifier for each DiaObject + - 'validityStartMjdTai' : Start time of object validity + + Returns + ------- + deduplication_map : `pandas.DataFrame` + Mapping of removed DiaObject IDs to kept IDs. Contains + columns: + + - 'removedDiaObjectId' : IDs to be removed + - 'keptDiaObjectId' : IDs to keep + + Only includes entries for clusters with multiple objects. + """ + + grp = diaObjects.groupby('cluster_label') + dedup_map = [] + for drp_id_i, group in grp: + if len(group) > 1: + ids = group['diaObjectId'].tolist() + # keep the oldest one + min_age_idx = group['validityStartMjdTai'].argmin() + id_keep = ids.pop(min_age_idx) + for idi in ids: + dedup_map.append((idi, id_keep)) + + return pd.DataFrame(dedup_map, columns=['removedDiaObjectId', 'keptDiaObjectId']) + + def cluster(self, diaObjects): + """Use agglomerative clustering to identify duplicate groups. + + This method applies hierarchical agglomerative clustering to + DiaObjects based on their sky positions (RA, Dec). Objects + within the configured maximum distance are grouped into + clusters representing duplicate sets. + + Parameters + ---------- + diaObjects : `pandas.DataFrame` + DataFrame of DiaObjects to cluster. Must include columns: + + - 'ra' : Right ascension in degrees + - 'dec' : Declination in degrees + + Returns + ------- + cluster_labels : `pandas.Series` + Series of cluster labels (integers) indexed to match the + input DataFrame. Objects with the same label belong to the + same cluster and are considered duplicates. + + Notes + ----- + A k-nearest neighbors connectivity graph is used to improve + performance. However, it prevents objects from being grouped + together if they are not among each other's k nearest + neighbors, even if they fall within the desired distance + threshold. `radius_neighbors_graph with a distance threshold + would avoid this issue but has proven computationally + infeasible. + """ + + coords = diaObjects[['ra', 'dec']].values + connectivity = kneighbors_graph(coords, n_neighbors=self.config.nNeighborsConnectivity, + mode='connectivity', include_self='auto') + + clustering = AgglomerativeClustering(distance_threshold=self.config.maxClusteringDistance/3600., + connectivity=connectivity, + n_clusters=None).fit(coords) + + self.log.info(f"Clustered {len(clustering.labels_)} diaObjects into " + f"{len(np.unique(clustering.labels_))} clusters.") + + return pd.Series(clustering.labels_, index=diaObjects.index, + name="cluster_label") + + def count_duplicates(self, catalog, radius_arcsec=None): + """Count catalog objects with self-matches within given radius. + + This method counts how many objects in the catalog have at least + one other object within the search radius, providing a measure + of the duplicate population. + + Parameters + ---------- + catalog : `pandas.DataFrame` + Catalog of objects to check for duplicates. + radius_arcsec : `float`, optional + Search radius in arcseconds within which to consider + objects as duplicates. If None, defaults to half of + config.maxClusteringDistance. + + Returns + ------- + duplicate_count : `int` + Number of objects that have at least one neighbor within + the specified radius (excluding self-matches). + + Notes + ----- + The method uses Astropy's catalog matching with nthneighbor=2 + to find the nearest neighbor excluding the object itself. + Objects whose nearest neighbor is within radius_arcsec are + counted as having duplicates. + + This provides a quick diagnostic count and may not exactly + match the number of duplicates identified by the clustering + algorithm. + """ + + if radius_arcsec is None: + radius_arcsec = self.config.maxClusteringDistance/2 + + sc = SkyCoord(catalog['ra'], catalog['dec'], unit='deg') + idx, d2d, _ = sc.match_to_catalog_sky(sc, nthneighbor=2) + wmatch = d2d < radius_arcsec*u.arcsecond + return np.sum(wmatch) diff --git a/tests/test_deduplicateAllSkyDiaObjects.py b/tests/test_deduplicateAllSkyDiaObjects.py new file mode 100644 index 00000000..b88b9578 --- /dev/null +++ b/tests/test_deduplicateAllSkyDiaObjects.py @@ -0,0 +1,298 @@ +# This file is part of ap_association. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +import os +import astropy.units +import numpy as np +import pandas as pd +import tempfile +import unittest + +from lsst.ap.association import DeduplicateAllSkyDiaObjectsTask +from lsst.ap.association.utils import getMidpointFromTimespan, readSchemaFromApdb +from lsst.dax.apdb import Apdb, ApdbSql +from lsst.utils import getPackageDir +import lsst.utils.tests +from utils_tests import makeExposure, makeDiaObjects, makeDiaSources, makeDiaForcedSources, makeRegionTime + + +def _data_file_name(basename, module_name): + """Return path name of a data file. + + Parameters + ---------- + basename : `str` + Name of the file to add to the path string. + module_name : `str` + Name of lsst stack package environment variable. + + Returns + ------- + data_file_path : `str` + Full path of the file to load from the "data" directory in a given + repository. + """ + return os.path.join(getPackageDir(module_name), "data", basename) + + +class TestDeduplicateAllSkyDiaObjects(unittest.TestCase): + + def setUp(self): + # Create an instance of random generator with fixed seed. + rng = np.random.default_rng(1234) + + self.db_file_fd, self.db_file = tempfile.mkstemp( + dir=os.path.dirname(__file__)) + self.addCleanup(os.remove, self.db_file) + self.addCleanup(os.close, self.db_file_fd) + + self.apdbConfig = ApdbSql.init_database(db_url="sqlite:///" + self.db_file) + self.config_file = tempfile.NamedTemporaryFile() + self.addCleanup(self.config_file.close) + self.apdbConfig.save(self.config_file.name) + self.apdb = Apdb.from_config(self.apdbConfig) + self.schema = readSchemaFromApdb(self.apdb) + + self.exposure = makeExposure(False, False) + self.regionTime = makeRegionTime(exposure=self.exposure) + self.dateTime = getMidpointFromTimespan(self.regionTime.timespan) + + original_diaObjects = makeDiaObjects(50, self.exposure, rng) + + # create duplicates + dup_0_0 = {"ra": original_diaObjects.iloc[0]['ra'], + "dec": original_diaObjects.iloc[0]['dec'] + 0.2/3600, + "diaObjectId": 1000, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_0_0["%s_psfFluxNdata" % f] = 0 + dup_0_1 = {"ra": original_diaObjects.iloc[0]['ra'] + 0.1/3600, + "dec": original_diaObjects.iloc[0]['dec'] + 0.2/3600, + "diaObjectId": 1001, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_0_1["%s_psfFluxNdata" % f] = 0 + dup_0_2 = {"ra": original_diaObjects.iloc[0]['ra'], + "dec": original_diaObjects.iloc[0]['dec'] - 0.1/3600, + "diaObjectId": 1002, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_0_2["%s_psfFluxNdata" % f] = 0 + dup_0_3 = {"ra": original_diaObjects.iloc[0]['ra'] + 0.2/3600, + "dec": original_diaObjects.iloc[0]['dec'] - 0.1/3600, + "diaObjectId": 1003, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_0_3["%s_psfFluxNdata" % f] = 0 + + dup_1_0 = {"ra": original_diaObjects.iloc[1]['ra'], + "dec": original_diaObjects.iloc[1]['dec'] + 0.2/3600, + "diaObjectId": 2000, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_1_0["%s_psfFluxNdata" % f] = 0 + dup_1_1 = {"ra": original_diaObjects.iloc[1]['ra'] + 0.1/3600, + "dec": original_diaObjects.iloc[1]['dec'] + 0.2/3600, + "diaObjectId": 2001, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_1_1["%s_psfFluxNdata" % f] = 0 + dup_1_2 = {"ra": original_diaObjects.iloc[1]['ra'], + "dec": original_diaObjects.iloc[1]['dec'] - 0.1/3600, + "diaObjectId": 2002, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_1_2["%s_psfFluxNdata" % f] = 0 + + dup_2_0 = {"ra": original_diaObjects.iloc[2]['ra'], + "dec": original_diaObjects.iloc[2]['dec'] + 0.2/3600, + "diaObjectId": 3000, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + dup_2_0["%s_psfFluxNdata" % f] = 0 + + data = [dup_0_0, dup_0_1, dup_0_2, dup_0_3, + dup_1_0, dup_1_1, dup_1_2, + dup_2_0] + duplicate_diaObjects = pd.DataFrame(data=data).set_index("diaObjectId", drop=False) + + self.diaObjects = pd.concat([original_diaObjects, duplicate_diaObjects]) + + self.diaSources = makeDiaSources( + 100, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) + self.diaForcedSources = makeDiaForcedSources( + 200, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) + + # Store the test diaSources as though they were observed a month before + # the current exposure. + dateTime = self.regionTime.timespan.begin.tai - 30 * astropy.units.day + self.apdb.store(dateTime, + self.diaObjects, + self.diaSources, + self.diaForcedSources) + + # These columns are not in the DPDD, yet do appear in DiaSource.yaml. + # We don't need to check them against the default APDB schema. + self.ignoreColumns = ["band", "bboxSize", "isDipole", "flags"] + + def _makeConfig(self, **kwargs): + config = DeduplicateAllSkyDiaObjectsTask.ConfigClass() + config.apdb_config_url = self.config_file.name + config.earliestMidpointMjdTai = self.dateTime.mjd + config.update(**kwargs) + return config + + def testRun(self): + """Test the full run method for deduplication. + """ + diaConfig = self._makeConfig() + task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) + result = task.run() + + self.assertIsNotNone(result.diaObjectDeduplicationMap) + + self.assertEqual(len(result.diaObjectDeduplicationMap), 8) + + self.assertIn('removedDiaObjectId', result.diaObjectDeduplicationMap.columns) + self.assertIn('keptDiaObjectId', result.diaObjectDeduplicationMap.columns) + + removed_ids = set(result.diaObjectDeduplicationMap['removedDiaObjectId']) + expected_duplicate_ids = {1000, 1001, 1002, 1003, 2000, 2001, 2002, 3000} + self.assertEqual(removed_ids, expected_duplicate_ids) + + kept_ids = set(result.diaObjectDeduplicationMap['keptDiaObjectId']) + self.assertTrue(kept_ids.issubset({1, 2, 3})) + + # Check that diaSources have been remapped from removed diaObjectIds + # (We can't easily compare to the kept ids because non-duplicated ids + # are not included in the deduplication map.) + updated_diaSources = self.apdb.getDiaSources(region=self.regionTime.region, object_ids=None, + visit_time=self.dateTime, start_time=None) + source_object_ids = set(updated_diaSources['diaObjectId']) + self.assertTrue(source_object_ids.isdisjoint(removed_ids), + "DiaSources should not reference removed diaObjectIds") + + def test_remap_clusters(self): + """Test the cluster remapping logic. + """ + diaConfig = self._makeConfig() + task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) + + # Create a simple test DataFrame with cluster labels + test_data = pd.DataFrame({ + 'diaObjectId': [1, 2, 3, 4, 5, 6], + 'validityStartMjdTai': [100.0, 101.0, 102.0, 200.0, 201.0, 300.0], + 'cluster_label': [0, 0, 0, 1, 1, 2], + 'ra': [10.0, 10.0, 10.0, 20.0, 20.0, 30.0], + 'dec': [5.0, 5.0, 5.0, 10.0, 10.0, 15.0] + }) + + dedup_map = task.remap_clusters(test_data) + + self.assertEqual(len(dedup_map), 3) # 2 + 1 remappings + + # Check that objects 2 and 3 map to object 1 + obj2_mapping = dedup_map[dedup_map['removedDiaObjectId'] == 2] + obj3_mapping = dedup_map[dedup_map['removedDiaObjectId'] == 3] + self.assertEqual(len(obj2_mapping), 1) + self.assertEqual(len(obj3_mapping), 1) + self.assertEqual(obj2_mapping.iloc[0]['keptDiaObjectId'], 1) + self.assertEqual(obj3_mapping.iloc[0]['keptDiaObjectId'], 1) + + # Check that object 5 maps to object 4 + obj5_mapping = dedup_map[dedup_map['removedDiaObjectId'] == 5] + self.assertEqual(len(obj5_mapping), 1) + self.assertEqual(obj5_mapping.iloc[0]['keptDiaObjectId'], 4) + + def test_cluster(self): + """Test the clustering algorithm. + """ + diaConfig = self._makeConfig(maxClusteringDistance=1.5) + task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) + + cluster_labels = task.cluster(self.diaObjects) + + self.assertEqual(len(cluster_labels), len(self.diaObjects)) + + group1_ids = [1, 1000, 1001, 1002, 1003] + group1_mask = self.diaObjects['diaObjectId'].isin(group1_ids) + group1_labels = cluster_labels[group1_mask] + + self.assertEqual(len(group1_labels.unique()), 1, + "Objects in duplicate group 1 should share a cluster label") + + group2_ids = [2, 2000, 2001, 2002] + group2_mask = self.diaObjects['diaObjectId'].isin(group2_ids) + group2_labels = cluster_labels[group2_mask] + + self.assertEqual(len(group2_labels.unique()), 1, + "Objects in duplicate group 2 should share a cluster label") + + group3_ids = [3, 3000] + group3_mask = self.diaObjects['diaObjectId'].isin(group3_ids) + group3_labels = cluster_labels[group3_mask] + + self.assertEqual(len(group3_labels.unique()), 1, + "Objects in duplicate group 3 should share a cluster label") + + self.assertNotEqual(group1_labels.iloc[0], group2_labels.iloc[0], + "Group 1 and Group 2 should have different cluster labels") + self.assertNotEqual(group1_labels.iloc[0], group3_labels.iloc[0], + "Group 1 and Group 3 should have different cluster labels") + self.assertNotEqual(group2_labels.iloc[0], group3_labels.iloc[0], + "Group 2 and Group 3 should have different cluster labels") + + def test_count_duplicates(self): + """Test the duplicate counting method. + """ + diaConfig = self._makeConfig(maxClusteringDistance=1.5) + task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) + + test_catalog = pd.DataFrame({ + 'diaObjectId': [1, 2, 3, 4, 5], + 'ra': [10.0, 10.0001, # Pair 1 + 10.1, 10.1001, # Pair 2 + 10.3], # Single + 'dec': [5.0, 5.0001, + 5.1, 5.1001, + 5.3] + }) + + duplicate_count = task.count_duplicates(test_catalog) + + self.assertEqual(duplicate_count, 4) + + # Test with custom radius + duplicate_count_small = task.count_duplicates(test_catalog, radius_arcsec=0.1) + self.assertEqual(duplicate_count_small, 0) + + # Test with large radius + duplicate_count_large = task.count_duplicates(test_catalog, radius_arcsec=3600.0) + self.assertEqual(duplicate_count_large, 5) + + +def setup_module(module): + lsst.utils.tests.init() + + +if __name__ == "__main__": + lsst.utils.tests.init() + unittest.main() From b38e97014d6fab86815ea55b1cb379a6f89314fc Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Thu, 29 Jan 2026 16:05:53 -0800 Subject: [PATCH 2/9] use new structure for reassignDiaSourcesToDiaObjects id_map --- python/lsst/ap/association/deduplicateAllSkyDiaObjects.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py index d2db91b6..5c40ae69 100644 --- a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py +++ b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py @@ -216,9 +216,13 @@ def remove_apdb_duplicates(self, diaObjects, deduplication_map): reassign_count = np.sum(w) if reassign_count: + wnew = diaObjects['diaObjectId'] == new_id + assert (np.sum(wnew) == 1) + newDiaObjectId = [daxApdb.recordIds.DiaObjectId.from_named_tuple(row) + for row in diaObjects.loc[wnew].itertuples()][0] for row in diaSourcesToReassign.loc[w].itertuples(): ds = daxApdb.recordIds.DiaSourceId.from_named_tuple(row) - id_map[ds] = new_id + id_map[ds] = newDiaObjectId self.log.verbose('Reassigned %d diaSources from diaObject %d to %d' % (reassign_count, old_id, new_id)) From 5c19bd86b3b32e46a798dce29dad1681078e8f06 Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Wed, 4 Feb 2026 09:37:13 -0800 Subject: [PATCH 3/9] Raise NoWorkFound if no DiaObjects are available. --- python/lsst/ap/association/deduplicateAllSkyDiaObjects.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py index 5c40ae69..1264f7b4 100644 --- a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py +++ b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py @@ -142,6 +142,9 @@ def run(self): diaObjects = all_diaObjects.loc[~wdup, :].copy() self.log.info(f"Filtered to {len(all_diaObjects)} latest DiaObjects.") + if len(diaObjects) == 0: + raise pipeBase.NoWorkFound("No DiaObjects available for deduplication.") + duplicate_count = self.count_duplicates(diaObjects) if duplicate_count == 0: From bfb9833191ac4582191391bde4edf260dbf1679f Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Mon, 2 Mar 2026 11:00:52 -0800 Subject: [PATCH 4/9] break up deduplication into spatial subsets --- .../deduplicateAllSkyDiaObjects.py | 154 +++++++++++++++++- 1 file changed, 148 insertions(+), 6 deletions(-) diff --git a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py index 1264f7b4..dd480044 100644 --- a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py +++ b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py @@ -25,8 +25,8 @@ import astropy.units as u import numpy as np import pandas as pd -from sklearn.cluster import AgglomerativeClustering -from sklearn.neighbors import kneighbors_graph +from sklearn.cluster import AgglomerativeClustering, MiniBatchKMeans +from sklearn.neighbors import kneighbors_graph, BallTree import lsst.dax.apdb as daxApdb import lsst.pex.config as pexConfig @@ -85,6 +85,14 @@ class DeduplicateAllSkyDiaObjectsConfig(pexConfig.Config): default=30, min=3, ) + maxSubsetSize = pexConfig.RangeField( + doc="Maximum number of DiaObjects to process in a single subset during " + "clustering. Larger values use more memory but may produce better " + "results at cluster boundaries.", + dtype=int, + default=50000, + min=1000, + ) class DeduplicateAllSkyDiaObjectsTask(pipeBase.Task): @@ -289,6 +297,12 @@ def cluster(self, diaObjects): within the configured maximum distance are grouped into clusters representing duplicate sets. + For large datasets, the data is first partitioned using KMeans + to create manageable subsets, then agglomerative clustering is + applied within each subset. Objects near partition boundaries + are processed with overlap to avoid missing cross-boundary + duplicates. + Parameters ---------- diaObjects : `pandas.DataFrame` @@ -314,18 +328,146 @@ def cluster(self, diaObjects): would avoid this issue but has proven computationally infeasible. """ + n_objects = len(diaObjects) + + # If the dataset is small enough, use direct clustering + if n_objects <= self.config.maxSubsetSize: + return self._cluster_subset(diaObjects) + # For large datasets, partition using KMeans first coords = diaObjects[['ra', 'dec']].values - connectivity = kneighbors_graph(coords, n_neighbors=self.config.nNeighborsConnectivity, + + # Determine number of partitions + n_partitions = max(2, int(np.ceil(n_objects / self.config.maxSubsetSize))) + self.log.info(f"Partitioning {n_objects} objects into {n_partitions} subsets " + "for memory-efficient clustering.") + + # Use MiniBatchKMeans for memory efficiency + kmeans = MiniBatchKMeans(n_clusters=n_partitions, random_state=42, + batch_size=min(10000, n_objects), + n_init=3) + partition_labels = kmeans.fit_predict(coords) + + # Initialize cluster labels with unique values per object + # We'll use a union-find approach to merge clusters across partitions + cluster_labels = np.arange(n_objects) + + # Build a BallTree for efficient neighbor queries + # Convert threshold to radians for haversine metric + threshold_rad = np.radians(self.config.maxClusteringDistance / 3600.0) + + # Process each partition with overlap + for partition_id in range(n_partitions): + # Get objects in this partition + partition_mask = partition_labels == partition_id + partition_indices = np.where(partition_mask)[0] + + if len(partition_indices) == 0: + continue + + # Find objects near partition boundaries by checking distance + # to objects in other partitions + partition_coords = coords[partition_indices] + + # Find nearby objects from other partitions that could be + # duplicates of objects in this partition + other_mask = ~partition_mask + other_indices = np.where(other_mask)[0] + +# if len(other_indices) > 0: + if False: + other_coords = coords[other_indices] + + # Use BallTree to find objects within threshold distance + tree = BallTree(np.radians(partition_coords), metric='haversine') + # Query which other objects are within threshold of + # partition objects + nearby_indices_list = tree.query_radius(np.radians(other_coords), + r=threshold_rad) + + # Find which other objects have neighbors in this partition + nearby_other_mask = np.array([len(x) > 0 for x in nearby_indices_list]) + nearby_other_indices = other_indices[nearby_other_mask] + + # Combine partition objects with nearby boundary objects + combined_indices = np.concatenate([partition_indices, nearby_other_indices]) + else: + combined_indices = partition_indices + + if len(combined_indices) == 0: + continue + + # Create subset DataFrame for clustering + subset_df = diaObjects.iloc[combined_indices] + + # Cluster this subset + subset_labels = self._cluster_subset(subset_df) + + # Map subset cluster labels back to global indices. + # Objects with the same subset label should have the same + # global label. + subset_indices = combined_indices + label_to_global = {} + + for i, (idx, label) in enumerate(zip(subset_indices, subset_labels.values)): + if label not in label_to_global: + # Use the first index we see for this label as the + # canonical one + label_to_global[label] = cluster_labels[idx] + else: + # Merge: set this object's label to match the canonical one + old_label = cluster_labels[idx] + new_label = label_to_global[label] + if old_label != new_label: + # Update all objects with old_label to new_label + cluster_labels[cluster_labels == old_label] = new_label + + # Renumber labels to be contiguous + unique_labels = np.unique(cluster_labels) + label_map = {old: new for new, old in enumerate(unique_labels)} + final_labels = np.array([label_map[lbl] for lbl in cluster_labels]) + + self.log.info(f"Clustered {n_objects} diaObjects into " + f"{len(unique_labels)} clusters.") + + return pd.Series(final_labels, index=diaObjects.index, + name="cluster_label") + + def _cluster_subset(self, diaObjects): + """Apply agglomerative clustering to a subset of DiaObjects. + + This is the core clustering routine applied to subsets that fit + in memory. + + Parameters + ---------- + diaObjects : `pandas.DataFrame` + DataFrame of DiaObjects to cluster. Must include columns: + + - 'ra' : Right ascension in degrees + - 'dec' : Declination in degrees + + Returns + ------- + cluster_labels : `pandas.Series` + Series of cluster labels (integers) indexed to match the + input DataFrame. + """ + coords = diaObjects[['ra', 'dec']].values + + # Adjust n_neighbors if subset is smaller than configured value + n_neighbors = min(self.config.nNeighborsConnectivity, len(diaObjects) - 1) + if n_neighbors < 1: + # Single object, just return label 0 + return pd.Series([0], index=diaObjects.index, name="cluster_label") + + connectivity = kneighbors_graph(coords, n_neighbors=n_neighbors, mode='connectivity', include_self='auto') clustering = AgglomerativeClustering(distance_threshold=self.config.maxClusteringDistance/3600., connectivity=connectivity, n_clusters=None).fit(coords) - self.log.info(f"Clustered {len(clustering.labels_)} diaObjects into " - f"{len(np.unique(clustering.labels_))} clusters.") - return pd.Series(clustering.labels_, index=diaObjects.index, name="cluster_label") From c44e91d38138258ba0006ce9a6517c9bbcbbd2d7 Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Wed, 20 May 2026 09:21:00 -0700 Subject: [PATCH 5/9] remove unneeded argument --- .../deduplicateAllSkyDiaObjects.py | 29 ++----------------- 1 file changed, 2 insertions(+), 27 deletions(-) diff --git a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py index dd480044..04be99d4 100644 --- a/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py +++ b/python/lsst/ap/association/deduplicateAllSkyDiaObjects.py @@ -37,19 +37,6 @@ class DeduplicateAllSkyDiaObjectsConfig(pexConfig.Config): """Configuration for DeduplicateAllSkyDiaObjectsTask. - - Parameters - ---------- - maxClusteringDistance : `float` - Maximum distance to merge clusters of duplicate DiaObjects, in - arcseconds. Default is 1.5 arcseconds. - earliestMidpointMjdTai : `float` - MidpointMjdTai (Modified Julian Date in TAI scale) of the - earliest DiaSource to be reassigned during deduplication. - Default corresponds to 2025-09-01. - maxDiaObjects : `int` - Maximum number of DiaObjects to process. An exception is raised - if more DiaObjects are found to prevent excessive memory/computation. """ apdb_config_url = pexConfig.Field( dtype=str, @@ -71,13 +58,6 @@ class DeduplicateAllSkyDiaObjectsConfig(pexConfig.Config): doc="MidpointMjdTai of earliest DIASource to be reassigned during " "deduplication.", ) - maxDiaObjects = pexConfig.RangeField( - doc="Maximum number of DiaObjects to process. An exception is raised " - "if more DiaObjects are found to prevent excessive memory/computation.", - dtype=int, - default=1000000, - min=1, - ) nNeighborsConnectivity = pexConfig.RangeField( doc="Number of neighbors to use for clustering. Larger values are more" "accurate but more expensive computationally.", @@ -137,10 +117,6 @@ def run(self): # fix getDiaObjectsForDedup returning a non-unique pandas index all_diaObjects = all_diaObjects.reset_index().drop('index', axis=1) - if len(all_diaObjects) > self.config.maxDiaObjects: - raise RuntimeError(f"Found {len(all_diaObjects)} DiaObjects, which exceeds the " - f"configured maximum of {self.config.maxDiaObjects}. " - "Aborting to avoid excessive memory/computation.") self.log.info(f"Loaded {len(all_diaObjects)} DiaObjects.") # only keep the latest versions of the DIAObjects @@ -324,7 +300,7 @@ def cluster(self, diaObjects): performance. However, it prevents objects from being grouped together if they are not among each other's k nearest neighbors, even if they fall within the desired distance - threshold. `radius_neighbors_graph with a distance threshold + threshold. `radius_neighbors_graph` with a distance threshold would avoid this issue but has proven computationally infeasible. """ @@ -374,8 +350,7 @@ def cluster(self, diaObjects): other_mask = ~partition_mask other_indices = np.where(other_mask)[0] -# if len(other_indices) > 0: - if False: + if len(other_indices) > 0: other_coords = coords[other_indices] # Use BallTree to find objects within threshold distance From 95362c4cc1beda42ac7810c6ac510f436b6690bb Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Wed, 20 May 2026 09:56:50 -0700 Subject: [PATCH 6/9] claude-authored unit tests for deduplication --- ...est_deduplicateAllSkyDiaObjects_cluster.py | 720 ++++++++++++++++++ 1 file changed, 720 insertions(+) create mode 100644 tests/test_deduplicateAllSkyDiaObjects_cluster.py diff --git a/tests/test_deduplicateAllSkyDiaObjects_cluster.py b/tests/test_deduplicateAllSkyDiaObjects_cluster.py new file mode 100644 index 00000000..b8fe450e --- /dev/null +++ b/tests/test_deduplicateAllSkyDiaObjects_cluster.py @@ -0,0 +1,720 @@ +# This file is part of ap_association. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . + +"""Comprehensive unit tests for the `cluster` method of +DeduplicateAllSkyDiaObjectsTask. +""" + +import numpy as np +import pandas as pd +import tempfile +import os +import unittest + +from lsst.ap.association import DeduplicateAllSkyDiaObjectsTask +from lsst.dax.apdb import ApdbSql +import lsst.utils.tests + + +def _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=30, + maxSubsetSize=50000): + """Create a DeduplicateAllSkyDiaObjectsTask with given config values. + + Creates a temporary SQLite APDB to satisfy the config requirement. + """ + db_file_fd, db_file = tempfile.mkstemp() + apdb_config = ApdbSql.init_database(db_url="sqlite:///" + db_file) + config_file = tempfile.NamedTemporaryFile(delete=False) + apdb_config.save(config_file.name) + config_file.close() + + config = DeduplicateAllSkyDiaObjectsTask.ConfigClass() + config.apdb_config_url = config_file.name + config.maxClusteringDistance = maxClusteringDistance + config.nNeighborsConnectivity = nNeighborsConnectivity + config.maxSubsetSize = maxSubsetSize + + task = DeduplicateAllSkyDiaObjectsTask(config=config) + + os.close(db_file_fd) + os.unlink(db_file) + os.unlink(config_file.name) + + return task + + +def _make_diaObjects(ra_list, dec_list, ids=None): + """Create a minimal DiaObjects DataFrame from lists of ra/dec.""" + n = len(ra_list) + if ids is None: + ids = list(range(1, n + 1)) + df = pd.DataFrame({ + 'diaObjectId': ids, + 'ra': ra_list, + 'dec': dec_list, + }) + df.index = pd.RangeIndex(n) + return df + + +class TestClusterSmallDataset(unittest.TestCase): + """Tests for the cluster method when dataset size <= maxSubsetSize + (direct clustering path via _cluster_subset). + """ + + def setUp(self): + self.task = _make_task(maxClusteringDistance=1.5) + + def test_single_object(self): + """A single object should get a unique cluster label.""" + diaObjects = _make_diaObjects([10.0], [5.0]) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 1) + self.assertEqual(labels.iloc[0], 0) + + def test_two_objects_within_threshold(self): + """Two objects within maxClusteringDistance should share a label.""" + # 0.3 arcsec apart in dec + diaObjects = _make_diaObjects( + [10.0, 10.0], + [5.0, 5.0 + 0.3 / 3600] + ) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 2) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_two_objects_beyond_threshold(self): + """Two objects beyond maxClusteringDistance should get different labels.""" + # 10 arcsec apart, well beyond the 1.5 arcsec threshold + diaObjects = _make_diaObjects( + [10.0, 10.0], + [5.0, 5.0 + 10.0 / 3600] + ) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 2) + self.assertNotEqual(labels.iloc[0], labels.iloc[1]) + + def test_multiple_distinct_clusters(self): + """Multiple well-separated groups should form distinct clusters.""" + # Three tight groups separated by large distances + ra = [10.0, 10.0 + 0.1 / 3600, # Group 1 + 20.0, 20.0 + 0.2 / 3600, # Group 2 + 30.0, 30.0 + 0.1 / 3600, 30.0 + 0.2 / 3600] # Group 3 + dec = [5.0, 5.0, + 15.0, 15.0, + 25.0, 25.0, 25.0] + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 7) + + # Group 1 should share a label + self.assertEqual(labels.iloc[0], labels.iloc[1]) + # Group 2 should share a label + self.assertEqual(labels.iloc[2], labels.iloc[3]) + # Group 3 should share a label + self.assertEqual(labels.iloc[4], labels.iloc[5]) + self.assertEqual(labels.iloc[5], labels.iloc[6]) + + # Distinct groups should have different labels + unique_group_labels = {labels.iloc[0], labels.iloc[2], labels.iloc[4]} + self.assertEqual(len(unique_group_labels), 3) + + def test_chain_of_objects(self): + """Objects forming a chain: each consecutive pair is within threshold, + but endpoints exceed threshold. + + AgglomerativeClustering with a distance_threshold will merge only + pairs within that distance. With a k-NN connectivity graph, it can + still link objects transitively through neighbors. However, the + distance_threshold limits merging to the configured distance, so + endpoints farther apart than the threshold may end up in separate + clusters depending on the linkage. + """ + # Create a chain of 3 objects, each 0.5 arcsec from its neighbor. + # Total span is 1.0 arcsec < threshold of 1.5 arcsec. + n = 3 + step = 0.5 / 3600 # 0.5 arcsec steps in degrees + ra = [10.0] * n + dec = [5.0 + i * step for i in range(n)] + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # All 3 are within 1.0 arcsec of each other, within the 1.5 arcsec + # threshold, so they should form a single cluster. + unique_labels = labels.unique() + self.assertEqual(len(unique_labels), 1, + "Chain of nearby objects within threshold should form a single cluster") + + def test_objects_exactly_at_threshold(self): + """Objects separated by exactly maxClusteringDistance are merged. + + AgglomerativeClustering uses <= for the distance_threshold, so + objects at exactly the threshold distance are merged. + """ + # Exactly 1.5 arcsec apart + diaObjects = _make_diaObjects( + [10.0, 10.0], + [5.0, 5.0 + 1.5 / 3600] + ) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 2) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_returns_series_with_matching_index(self): + """The returned Series should have the same index as the input.""" + diaObjects = _make_diaObjects([10.0, 20.0, 30.0], [5.0, 15.0, 25.0]) + diaObjects.index = pd.Index([10, 20, 30]) + + labels = self.task.cluster(diaObjects) + + self.assertIsInstance(labels, pd.Series) + self.assertTrue(labels.index.equals(diaObjects.index)) + self.assertEqual(labels.name, "cluster_label") + + def test_contiguous_labels(self): + """Cluster labels should be contiguous integers starting from 0.""" + # Create 10 well-separated objects + ra = [float(i * 10) for i in range(10)] + dec = [float(i * 10) for i in range(10)] + diaObjects = _make_diaObjects(ra, dec) + + labels = self.task.cluster(diaObjects) + + unique_labels = sorted(labels.unique()) + self.assertEqual(unique_labels, list(range(len(unique_labels)))) + + def test_large_cluster(self): + """A tight group of many objects should mostly cluster together. + + With k-NN connectivity, very tight groups where all objects are + mutual nearest neighbors should form a single cluster. We use + a small group here to ensure the connectivity graph is dense enough. + """ + rng = np.random.default_rng(42) + n = 15 + # Objects within 0.3 arcsec of center (well within 1.5 arcsec threshold) + center_ra, center_dec = 100.0, -30.0 + ra = center_ra + rng.uniform(-0.3 / 3600, 0.3 / 3600, n) + dec = center_dec + rng.uniform(-0.3 / 3600, 0.3 / 3600, n) + + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + labels = self.task.cluster(diaObjects) + + # All should be in the same cluster + self.assertEqual(len(labels.unique()), 1) + + def test_no_duplicates_all_isolated(self): + """When all objects are well separated, each gets its own cluster.""" + n = 20 + # Place objects 1 degree apart + ra = [float(i) for i in range(n)] + dec = [0.0] * n + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), n) + + def test_mixed_duplicates_and_isolated(self): + """A mix of duplicate groups and isolated objects.""" + ra = [10.0, 10.0 + 0.1 / 3600, 10.0 + 0.2 / 3600, # Group of 3 + 50.0, # Isolated + 80.0, 80.0 + 0.1 / 3600, # Pair + 120.0] # Isolated + dec = [5.0, 5.0, 5.0, + 20.0, + 40.0, 40.0, + 60.0] + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # Group of 3 should share a label + self.assertEqual(labels.iloc[0], labels.iloc[1]) + self.assertEqual(labels.iloc[1], labels.iloc[2]) + + # Pair should share a label + self.assertEqual(labels.iloc[4], labels.iloc[5]) + + # All distinct clusters + unique_groups = {labels.iloc[0], labels.iloc[3], labels.iloc[4], labels.iloc[6]} + self.assertEqual(len(unique_groups), 4) + + def test_nNeighborsConnectivity_adapted_for_small_n(self): + """When n_objects < nNeighborsConnectivity, the method should adapt.""" + # config nNeighborsConnectivity=30 but only 5 objects + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 20.0, 30.0, 40.0], + [5.0, 5.0, 15.0, 25.0, 35.0] + ) + # Should not raise + labels = self.task.cluster(diaObjects) + self.assertEqual(len(labels), 5) + # First two should be clustered + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_custom_clustering_distance(self): + """Changing maxClusteringDistance affects cluster membership.""" + ra = [10.0, 10.0] + dec = [5.0, 5.0 + 1.0 / 3600] # 1 arcsec apart + + diaObjects = _make_diaObjects(ra, dec) + + # With 1.5 arcsec threshold, should cluster + task_wide = _make_task(maxClusteringDistance=1.5) + labels_wide = task_wide.cluster(diaObjects) + self.assertEqual(labels_wide.iloc[0], labels_wide.iloc[1]) + + # With 0.5 arcsec threshold, should not cluster + task_narrow = _make_task(maxClusteringDistance=0.5) + labels_narrow = task_narrow.cluster(diaObjects) + self.assertNotEqual(labels_narrow.iloc[0], labels_narrow.iloc[1]) + + def test_declination_near_pole(self): + """Objects near the celestial pole should still cluster correctly.""" + # Near north pole + diaObjects = _make_diaObjects( + [0.0, 180.0], # Opposite RA but very close at the pole + [89.9999, 89.9999] + ) + labels = self.task.cluster(diaObjects) + # These are actually close together at the pole + # (cos(dec) factor makes RA separation small) + # but AgglomerativeClustering uses Euclidean on (ra, dec) coords, + # so it will treat them as 180 degrees apart in RA + self.assertEqual(len(labels), 2) + + def test_ra_wraparound(self): + """Objects near RA=0/360 boundary.""" + # Two objects near RA=0, one just above and one just below + diaObjects = _make_diaObjects( + [0.0001, 359.9999], + [5.0, 5.0] + ) + labels = self.task.cluster(diaObjects) + # Euclidean clustering in (ra, dec) space won't handle wraparound; + # these will appear as ~360 degrees apart and won't cluster. + self.assertNotEqual(labels.iloc[0], labels.iloc[1]) + + +class TestClusterLargeDataset(unittest.TestCase): + """Tests for the cluster method when dataset size > maxSubsetSize + (partitioned clustering path via KMeans + boundary merging). + """ + + def setUp(self): + # Use a very small maxSubsetSize to trigger partitioning + self.task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + + def test_partitioning_triggered(self): + """Verify that partitioning is used for large datasets.""" + rng = np.random.default_rng(42) + n = 2000 + ra = rng.uniform(0, 360, n) + dec = rng.uniform(-90, 90, n) + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), n) + # Should have many unique clusters (most objects are isolated) + self.assertGreater(len(labels.unique()), n * 0.9) + + def test_duplicates_found_within_partition(self): + """Duplicates within the same partition should be identified.""" + rng = np.random.default_rng(42) + n = 1500 # Exceeds maxSubsetSize=1000 + + # Place most objects randomly and far apart + ra = rng.uniform(0, 180, n).tolist() + dec = rng.uniform(-45, 45, n).tolist() + + # Add a tight duplicate group (will end up in same partition) + group_center_ra = 90.0 + group_center_dec = 0.0 + for i in range(5): + ra.append(group_center_ra + rng.uniform(-0.3, 0.3) / 3600) + dec.append(group_center_dec + rng.uniform(-0.3, 0.3) / 3600) + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # The 5 duplicate objects at the end should share a cluster label + dup_labels = labels.iloc[n:n + 5] + self.assertEqual(len(dup_labels.unique()), 1, + "Tight group within partition should be clustered") + + def test_duplicates_at_partition_boundary(self): + """Duplicates that span a partition boundary should still be merged.""" + # maxSubsetSize=1000, so with 1500 objects we get 2 partitions. + # Create two clusters of objects placed such that KMeans will + # split them into different partitions, with a duplicate pair + # at the boundary. + rng = np.random.default_rng(123) + + # Group A: objects in one region of sky + n_a = 800 + ra_a = rng.uniform(10, 20, n_a).tolist() + dec_a = rng.uniform(10, 20, n_a).tolist() + + # Group B: objects in a different region + n_b = 800 + ra_b = rng.uniform(100, 110, n_b).tolist() + dec_b = rng.uniform(-20, -10, n_b).tolist() + + # Now add a pair of duplicates at the midpoint between groups + # These are close together but KMeans may split them + mid_ra = 55.0 + mid_dec = 0.0 + ra_boundary = [mid_ra, mid_ra + 0.2 / 3600] + dec_boundary = [mid_dec, mid_dec + 0.1 / 3600] + + ra = ra_a + ra_b + ra_boundary + dec = dec_a + dec_b + dec_boundary + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # The boundary pair should be clustered together + boundary_labels = labels.iloc[n_a + n_b:] + self.assertEqual(boundary_labels.iloc[0], boundary_labels.iloc[1], + "Boundary duplicate pair should share a cluster label") + + def test_output_labels_are_contiguous(self): + """Labels should be renumbered to be contiguous even after merging.""" + rng = np.random.default_rng(99) + n = 1500 + ra = rng.uniform(0, 360, n) + dec = rng.uniform(-90, 90, n) + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + + labels = self.task.cluster(diaObjects) + + unique_labels = sorted(labels.unique()) + expected = list(range(len(unique_labels))) + self.assertEqual(unique_labels, expected) + + def test_returns_series_with_matching_index(self): + """Returned Series matches input DataFrame index.""" + rng = np.random.default_rng(77) + n = 1500 + ra = rng.uniform(0, 360, n) + dec = rng.uniform(-90, 90, n) + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + diaObjects.index = pd.RangeIndex(start=100, stop=100 + n) + + labels = self.task.cluster(diaObjects) + + self.assertIsInstance(labels, pd.Series) + self.assertTrue(labels.index.equals(diaObjects.index)) + self.assertEqual(labels.name, "cluster_label") + + def test_consistency_with_small_path(self): + """Results from the partitioned path should be consistent with + the direct path for datasets that can use either. + """ + rng = np.random.default_rng(55) + n = 1200 + + # Create data with known duplicate structure + # 10 well-separated groups, each with 2-3 tight duplicates + ra = [] + dec = [] + for i in range(10): + center_ra = float(i * 30) + center_dec = float(i * 10 - 45) + n_in_group = rng.integers(2, 4) + for _ in range(n_in_group): + ra.append(center_ra + rng.uniform(-0.3, 0.3) / 3600) + dec.append(center_dec + rng.uniform(-0.3, 0.3) / 3600) + + # Fill remaining with isolated objects + n_placed = len(ra) + for i in range(n - n_placed): + ra.append(rng.uniform(0, 360)) + dec.append(rng.uniform(-90, 90)) + + diaObjects = _make_diaObjects(ra, dec) + + # Run with small maxSubsetSize (partitioned path) + task_partitioned = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + labels_partitioned = task_partitioned.cluster(diaObjects) + + # Run with large maxSubsetSize (direct path) + task_direct = _make_task(maxClusteringDistance=1.5, maxSubsetSize=50000) + labels_direct = task_direct.cluster(diaObjects) + + # The known tight groups should be clustered in both cases + idx = 0 + for i in range(10): + n_in_group = 0 + group_start = idx + center_ra = float(i * 30) + center_dec = float(i * 10 - 45) + while idx < n_placed: + if (abs(ra[idx] - center_ra) < 1.0 / 3600 + and abs(dec[idx] - center_dec) < 1.0 / 3600): + n_in_group += 1 + idx += 1 + else: + break + + if n_in_group > 1: + # Check direct path clusters them + direct_group = labels_direct.iloc[group_start:group_start + n_in_group] + self.assertEqual(len(direct_group.unique()), 1, + f"Direct path: group {i} should be one cluster") + + # Check partitioned path clusters them + part_group = labels_partitioned.iloc[group_start:group_start + n_in_group] + self.assertEqual(len(part_group.unique()), 1, + f"Partitioned path: group {i} should be one cluster") + + def test_empty_partition_handled(self): + """An empty KMeans partition should be handled gracefully.""" + # We can't easily force an empty partition from KMeans, but we can + # verify the code runs without error on a dataset where some + # partitions might be very small. + rng = np.random.default_rng(7) + # Highly clustered data that may produce uneven partitions + n = 1500 + # Most objects in one location + ra = rng.normal(50.0, 0.01, n).tolist() + dec = rng.normal(20.0, 0.01, n).tolist() + # A few outliers + for i in range(50): + ra.append(float(i * 7)) + dec.append(float(i * 3 - 45)) + + diaObjects = _make_diaObjects(ra + ra[-50:], dec + dec[-50:]) + # Should complete without error + labels = self.task.cluster(diaObjects) + self.assertEqual(len(labels), len(diaObjects)) + + def test_many_partitions(self): + """A dataset requiring many partitions.""" + task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + rng = np.random.default_rng(2024) + n = 5000 # Will create ~5 partitions + + ra = rng.uniform(0, 360, n).tolist() + dec = rng.uniform(-90, 90, n).tolist() + + # Add known duplicates + for i in range(20): + base_ra = rng.uniform(0, 360) + base_dec = rng.uniform(-60, 60) + ra.extend([base_ra, base_ra + 0.2 / 3600]) + dec.extend([base_dec, base_dec + 0.1 / 3600]) + + diaObjects = _make_diaObjects(ra, dec) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels), len(diaObjects)) + + # Verify known duplicates are clustered + for i in range(20): + idx1 = n + 2 * i + idx2 = n + 2 * i + 1 + self.assertEqual(labels.iloc[idx1], labels.iloc[idx2], + f"Known duplicate pair {i} should be clustered") + + +class TestClusterSubset(unittest.TestCase): + """Direct tests of the _cluster_subset helper.""" + + def setUp(self): + self.task = _make_task(maxClusteringDistance=1.5) + + def test_single_object_returns_label_zero(self): + """A single object returns label 0.""" + diaObjects = _make_diaObjects([10.0], [5.0]) + labels = self.task._cluster_subset(diaObjects) + + self.assertEqual(len(labels), 1) + self.assertEqual(labels.iloc[0], 0) + + def test_two_nearby_objects(self): + """Two nearby objects get the same label.""" + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600], + [5.0, 5.0] + ) + labels = self.task._cluster_subset(diaObjects) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_preserves_dataframe_index(self): + """Labels index should match input DataFrame index.""" + diaObjects = _make_diaObjects([10.0, 20.0, 30.0], [5.0, 15.0, 25.0]) + diaObjects.index = pd.Index([5, 10, 15]) + + labels = self.task._cluster_subset(diaObjects) + self.assertTrue(labels.index.equals(diaObjects.index)) + + +class TestClusterEdgeCases(unittest.TestCase): + """Edge cases and special configurations for the cluster method.""" + + def test_all_objects_identical_position(self): + """All objects at the same location should form one cluster.""" + n = 10 + ra = [100.0] * n + dec = [-30.0] * n + diaObjects = _make_diaObjects(ra, dec) + + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), 1) + + def test_negative_declination(self): + """Clustering works for objects in the southern hemisphere.""" + diaObjects = _make_diaObjects( + [200.0, 200.0 + 0.1 / 3600, 200.0 + 0.2 / 3600], + [-60.0, -60.0, -60.0] + ) + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), 1) + + def test_objects_along_equator(self): + """Objects along the celestial equator.""" + # Objects 0.5 arcsec apart along RA at dec=0 + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.5 / 3600, 10.0 + 1.0 / 3600], + [0.0, 0.0, 0.0] + ) + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + # All three within 1 arcsec of each other, should cluster + self.assertEqual(len(labels.unique()), 1) + + def test_minimum_nNeighborsConnectivity(self): + """Task with minimum nNeighborsConnectivity=3 still works.""" + task = _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=3) + + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 20.0, 30.0, 40.0], + [5.0, 5.0, 15.0, 25.0, 35.0] + ) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels), 5) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_maxSubsetSize_boundary(self): + """Dataset exactly at maxSubsetSize should use direct path.""" + n = 1000 + task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + + rng = np.random.default_rng(42) + ra = rng.uniform(0, 360, n).tolist() + dec = rng.uniform(-90, 90, n).tolist() + diaObjects = _make_diaObjects(ra, dec) + + # Should use direct path (n == maxSubsetSize) + labels = task.cluster(diaObjects) + self.assertEqual(len(labels), n) + + def test_maxSubsetSize_plus_one_triggers_partitioning(self): + """Dataset of maxSubsetSize+1 should trigger partitioning.""" + n = 1001 + task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + + rng = np.random.default_rng(42) + ra = rng.uniform(0, 360, n).tolist() + dec = rng.uniform(-90, 90, n).tolist() + diaObjects = _make_diaObjects(ra, dec) + + labels = task.cluster(diaObjects) + self.assertEqual(len(labels), n) + + def test_three_objects_same_position(self): + """Three objects at the same position form one cluster.""" + diaObjects = _make_diaObjects( + [45.0, 45.0, 45.0], + [-10.0, -10.0, -10.0] + ) + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), 1) + + def test_non_default_index(self): + """DataFrame with a non-default integer index.""" + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 50.0], + [5.0, 5.0, 30.0] + ) + diaObjects.index = pd.Index([100, 200, 300]) + + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertTrue(labels.index.equals(diaObjects.index)) + self.assertEqual(labels.loc[100], labels.loc[200]) + self.assertNotEqual(labels.loc[100], labels.loc[300]) + + def test_large_nNeighborsConnectivity(self): + """nNeighborsConnectivity larger than dataset is handled.""" + task = _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=30) + + # Only 4 objects but nNeighborsConnectivity=30 + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 50.0, 80.0], + [5.0, 5.0, 30.0, 60.0] + ) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels), 4) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_deterministic_results(self): + """Calling cluster twice on the same data yields identical results.""" + rng = np.random.default_rng(42) + n = 100 + ra = rng.uniform(0, 360, n).tolist() + dec = rng.uniform(-90, 90, n).tolist() + diaObjects = _make_diaObjects(ra, dec) + + task = _make_task(maxClusteringDistance=1.5) + labels1 = task.cluster(diaObjects) + labels2 = task.cluster(diaObjects) + + pd.testing.assert_series_equal(labels1, labels2) + + +def setup_module(module): + lsst.utils.tests.init() + + +if __name__ == "__main__": + lsst.utils.tests.init() + unittest.main() From b8cc910b43282e2e4904506cc8477501809636cc Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Wed, 20 May 2026 14:02:53 -0700 Subject: [PATCH 7/9] merge unit tests --- tests/test_deduplicateAllSkyDiaObjects.py | 655 ++++++++++++++++ ...est_deduplicateAllSkyDiaObjects_cluster.py | 720 ------------------ 2 files changed, 655 insertions(+), 720 deletions(-) delete mode 100644 tests/test_deduplicateAllSkyDiaObjects_cluster.py diff --git a/tests/test_deduplicateAllSkyDiaObjects.py b/tests/test_deduplicateAllSkyDiaObjects.py index b88b9578..455b7e32 100644 --- a/tests/test_deduplicateAllSkyDiaObjects.py +++ b/tests/test_deduplicateAllSkyDiaObjects.py @@ -289,6 +289,661 @@ def test_count_duplicates(self): self.assertEqual(duplicate_count_large, 5) + + +def _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=30, + maxSubsetSize=50000): + """Create a DeduplicateAllSkyDiaObjectsTask with given config values. + + Creates a temporary SQLite APDB to satisfy the config requirement. + """ + db_file_fd, db_file = tempfile.mkstemp() + apdb_config = ApdbSql.init_database(db_url="sqlite:///" + db_file) + config_file = tempfile.NamedTemporaryFile(delete=False) + apdb_config.save(config_file.name) + config_file.close() + + config = DeduplicateAllSkyDiaObjectsTask.ConfigClass() + config.apdb_config_url = config_file.name + config.maxClusteringDistance = maxClusteringDistance + config.nNeighborsConnectivity = nNeighborsConnectivity + config.maxSubsetSize = maxSubsetSize + + task = DeduplicateAllSkyDiaObjectsTask(config=config) + + os.close(db_file_fd) + os.unlink(db_file) + os.unlink(config_file.name) + + return task + + +def _make_diaObjects(ra_list, dec_list, ids=None): + """Create a minimal DiaObjects DataFrame from lists of ra/dec.""" + n = len(ra_list) + if ids is None: + ids = list(range(1, n + 1)) + df = pd.DataFrame({ + 'diaObjectId': ids, + 'ra': ra_list, + 'dec': dec_list, + }) + df.index = pd.RangeIndex(n) + return df + + +class TestClusterSmallDataset(unittest.TestCase): + """Tests for the cluster method when dataset size <= maxSubsetSize + (direct clustering path via _cluster_subset). + """ + + def setUp(self): + self.task = _make_task(maxClusteringDistance=1.5) + + def test_single_object(self): + """A single object should get a unique cluster label.""" + diaObjects = _make_diaObjects([10.0], [5.0]) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 1) + self.assertEqual(labels.iloc[0], 0) + + def test_two_objects_within_threshold(self): + """Two objects within maxClusteringDistance should share a label.""" + # 0.3 arcsec apart in dec + diaObjects = _make_diaObjects( + [10.0, 10.0], + [5.0, 5.0 + 0.3 / 3600] + ) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 2) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_two_objects_beyond_threshold(self): + """Two objects beyond maxClusteringDistance should get different labels.""" + # 10 arcsec apart, well beyond the 1.5 arcsec threshold + diaObjects = _make_diaObjects( + [10.0, 10.0], + [5.0, 5.0 + 10.0 / 3600] + ) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 2) + self.assertNotEqual(labels.iloc[0], labels.iloc[1]) + + def test_multiple_distinct_clusters(self): + """Multiple well-separated groups should form distinct clusters.""" + # Three tight groups separated by large distances + ra = [10.0, 10.0 + 0.1 / 3600, # Group 1 + 20.0, 20.0 + 0.2 / 3600, # Group 2 + 30.0, 30.0 + 0.1 / 3600, 30.0 + 0.2 / 3600] # Group 3 + dec = [5.0, 5.0, + 15.0, 15.0, + 25.0, 25.0, 25.0] + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 7) + + # Group 1 should share a label + self.assertEqual(labels.iloc[0], labels.iloc[1]) + # Group 2 should share a label + self.assertEqual(labels.iloc[2], labels.iloc[3]) + # Group 3 should share a label + self.assertEqual(labels.iloc[4], labels.iloc[5]) + self.assertEqual(labels.iloc[5], labels.iloc[6]) + + # Distinct groups should have different labels + unique_group_labels = {labels.iloc[0], labels.iloc[2], labels.iloc[4]} + self.assertEqual(len(unique_group_labels), 3) + + def test_chain_of_objects(self): + """Objects forming a chain: each consecutive pair is within threshold, + but endpoints exceed threshold. + + AgglomerativeClustering with a distance_threshold will merge only + pairs within that distance. With a k-NN connectivity graph, it can + still link objects transitively through neighbors. However, the + distance_threshold limits merging to the configured distance, so + endpoints farther apart than the threshold may end up in separate + clusters depending on the linkage. + """ + # Create a chain of 3 objects, each 0.5 arcsec from its neighbor. + # Total span is 1.0 arcsec < threshold of 1.5 arcsec. + n = 3 + step = 0.5 / 3600 # 0.5 arcsec steps in degrees + ra = [10.0] * n + dec = [5.0 + i * step for i in range(n)] + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # All 3 are within 1.0 arcsec of each other, within the 1.5 arcsec + # threshold, so they should form a single cluster. + unique_labels = labels.unique() + self.assertEqual(len(unique_labels), 1, + "Chain of nearby objects within threshold should form a single cluster") + + def test_objects_exactly_at_threshold(self): + """Objects separated by exactly maxClusteringDistance are merged. + + AgglomerativeClustering uses <= for the distance_threshold, so + objects at exactly the threshold distance are merged. + """ + # Exactly 1.5 arcsec apart + diaObjects = _make_diaObjects( + [10.0, 10.0], + [5.0, 5.0 + 1.5 / 3600] + ) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), 2) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_returns_series_with_matching_index(self): + """The returned Series should have the same index as the input.""" + diaObjects = _make_diaObjects([10.0, 20.0, 30.0], [5.0, 15.0, 25.0]) + diaObjects.index = pd.Index([10, 20, 30]) + + labels = self.task.cluster(diaObjects) + + self.assertIsInstance(labels, pd.Series) + self.assertTrue(labels.index.equals(diaObjects.index)) + self.assertEqual(labels.name, "cluster_label") + + + def test_large_cluster(self): + """A tight group of many objects should mostly cluster together. + + With k-NN connectivity, very tight groups where all objects are + mutual nearest neighbors should form a single cluster. We use + a small group here to ensure the connectivity graph is dense enough. + """ + rng = np.random.default_rng(42) + n = 15 + # Objects within 0.3 arcsec of center (well within 1.5 arcsec threshold) + center_ra, center_dec = 100.0, -30.0 + ra = center_ra + rng.uniform(-0.3 / 3600, 0.3 / 3600, n) + dec = center_dec + rng.uniform(-0.3 / 3600, 0.3 / 3600, n) + + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + labels = self.task.cluster(diaObjects) + + # All should be in the same cluster + self.assertEqual(len(labels.unique()), 1) + + def test_no_duplicates_all_isolated(self): + """When all objects are well separated, each gets its own cluster.""" + n = 20 + # Place objects 1 degree apart + ra = [float(i) for i in range(n)] + dec = [0.0] * n + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), n) + + def test_mixed_duplicates_and_isolated(self): + """A mix of duplicate groups and isolated objects.""" + ra = [10.0, 10.0 + 0.1 / 3600, 10.0 + 0.2 / 3600, # Group of 3 + 50.0, # Isolated + 80.0, 80.0 + 0.1 / 3600, # Pair + 120.0] # Isolated + dec = [5.0, 5.0, 5.0, + 20.0, + 40.0, 40.0, + 60.0] + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # Group of 3 should share a label + self.assertEqual(labels.iloc[0], labels.iloc[1]) + self.assertEqual(labels.iloc[1], labels.iloc[2]) + + # Pair should share a label + self.assertEqual(labels.iloc[4], labels.iloc[5]) + + # All distinct clusters + unique_groups = {labels.iloc[0], labels.iloc[3], labels.iloc[4], labels.iloc[6]} + self.assertEqual(len(unique_groups), 4) + + def test_nNeighborsConnectivity_adapted_for_small_n(self): + """When n_objects < nNeighborsConnectivity, the method should adapt.""" + # config nNeighborsConnectivity=30 but only 5 objects + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 20.0, 30.0, 40.0], + [5.0, 5.0, 15.0, 25.0, 35.0] + ) + # Should not raise + labels = self.task.cluster(diaObjects) + self.assertEqual(len(labels), 5) + # First two should be clustered + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_custom_clustering_distance(self): + """Changing maxClusteringDistance affects cluster membership.""" + ra = [10.0, 10.0] + dec = [5.0, 5.0 + 1.0 / 3600] # 1 arcsec apart + + diaObjects = _make_diaObjects(ra, dec) + + # With 1.5 arcsec threshold, should cluster + task_wide = _make_task(maxClusteringDistance=1.5) + labels_wide = task_wide.cluster(diaObjects) + self.assertEqual(labels_wide.iloc[0], labels_wide.iloc[1]) + + # With 0.5 arcsec threshold, should not cluster + task_narrow = _make_task(maxClusteringDistance=0.5) + labels_narrow = task_narrow.cluster(diaObjects) + self.assertNotEqual(labels_narrow.iloc[0], labels_narrow.iloc[1]) + + def test_declination_near_pole(self): + """Objects near the celestial pole should still cluster correctly.""" + # Near north pole + diaObjects = _make_diaObjects( + [0.0, 180.0], # Opposite RA but very close at the pole + [89.9999, 89.9999] + ) + labels = self.task.cluster(diaObjects) + # These are actually close together at the pole + # (cos(dec) factor makes RA separation small) + # but AgglomerativeClustering uses Euclidean on (ra, dec) coords, + # so it will treat them as 180 degrees apart in RA + self.assertEqual(len(labels), 2) + + def test_ra_wraparound(self): + """Objects near RA=0/360 boundary.""" + # Two objects near RA=0, one just above and one just below + diaObjects = _make_diaObjects( + [0.0001, 359.9999], + [5.0, 5.0] + ) + labels = self.task.cluster(diaObjects) + # Euclidean clustering in (ra, dec) space won't handle wraparound; + # these will appear as ~360 degrees apart and won't cluster. + self.assertNotEqual(labels.iloc[0], labels.iloc[1]) + + +class TestClusterLargeDataset(unittest.TestCase): + """Tests for the cluster method when dataset size > maxSubsetSize + (partitioned clustering path via KMeans + boundary merging). + """ + + def setUp(self): + # Use a very small maxSubsetSize to trigger partitioning + self.task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + + def test_partitioning_triggered(self): + """Verify that partitioning is used for large datasets.""" + rng = np.random.default_rng(42) + n = 2000 + ra = rng.uniform(0, 360, n) + dec = rng.uniform(-90, 90, n) + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + + labels = self.task.cluster(diaObjects) + + self.assertEqual(len(labels), n) + # Should have many unique clusters (most objects are isolated) + self.assertGreater(len(labels.unique()), n * 0.9) + + def test_duplicates_found_within_partition(self): + """Duplicates within the same partition should be identified.""" + rng = np.random.default_rng(42) + n = 1500 # Exceeds maxSubsetSize=1000 + + # Place most objects randomly and far apart + ra = rng.uniform(0, 180, n).tolist() + dec = rng.uniform(-45, 45, n).tolist() + + # Add a tight duplicate group (will end up in same partition) + group_center_ra = 90.0 + group_center_dec = 0.0 + for i in range(5): + ra.append(group_center_ra + rng.uniform(-0.3, 0.3) / 3600) + dec.append(group_center_dec + rng.uniform(-0.3, 0.3) / 3600) + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # The 5 duplicate objects at the end should share a cluster label + dup_labels = labels.iloc[n:n + 5] + self.assertEqual(len(dup_labels.unique()), 1, + "Tight group within partition should be clustered") + + def test_duplicates_at_partition_boundary(self): + """Duplicates that span a partition boundary should still be merged.""" + # maxSubsetSize=1000, so with 1500 objects we get 2 partitions. + # Create two clusters of objects placed such that KMeans will + # split them into different partitions, with a duplicate pair + # at the boundary. + rng = np.random.default_rng(123) + + # Group A: objects in one region of sky + n_a = 800 + ra_a = rng.uniform(10, 20, n_a).tolist() + dec_a = rng.uniform(10, 20, n_a).tolist() + + # Group B: objects in a different region + n_b = 800 + ra_b = rng.uniform(100, 110, n_b).tolist() + dec_b = rng.uniform(-20, -10, n_b).tolist() + + # Now add a pair of duplicates at the midpoint between groups + # These are close together but KMeans may split them + mid_ra = 55.0 + mid_dec = 0.0 + ra_boundary = [mid_ra, mid_ra + 0.2 / 3600] + dec_boundary = [mid_dec, mid_dec + 0.1 / 3600] + + ra = ra_a + ra_b + ra_boundary + dec = dec_a + dec_b + dec_boundary + + diaObjects = _make_diaObjects(ra, dec) + labels = self.task.cluster(diaObjects) + + # The boundary pair should be clustered together + boundary_labels = labels.iloc[n_a + n_b:] + self.assertEqual(boundary_labels.iloc[0], boundary_labels.iloc[1], + "Boundary duplicate pair should share a cluster label") + + def test_output_labels_are_contiguous(self): + """Labels should be renumbered to be contiguous even after merging.""" + rng = np.random.default_rng(99) + n = 1500 + ra = rng.uniform(0, 360, n) + dec = rng.uniform(-90, 90, n) + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + + labels = self.task.cluster(diaObjects) + + unique_labels = sorted(labels.unique()) + expected = list(range(len(unique_labels))) + self.assertEqual(unique_labels, expected) + + def test_returns_series_with_matching_index(self): + """Returned Series matches input DataFrame index.""" + rng = np.random.default_rng(77) + n = 1500 + ra = rng.uniform(0, 360, n) + dec = rng.uniform(-90, 90, n) + diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) + diaObjects.index = pd.RangeIndex(start=100, stop=100 + n) + + labels = self.task.cluster(diaObjects) + + self.assertIsInstance(labels, pd.Series) + self.assertTrue(labels.index.equals(diaObjects.index)) + self.assertEqual(labels.name, "cluster_label") + + def test_consistency_with_small_path(self): + """Results from the partitioned path should be consistent with + the direct path for datasets that can use either. + """ + rng = np.random.default_rng(55) + n = 1200 + + # Create data with known duplicate structure + # 10 well-separated groups, each with 2-3 tight duplicates + ra = [] + dec = [] + for i in range(10): + center_ra = float(i * 30) + center_dec = float(i * 10 - 45) + n_in_group = rng.integers(2, 4) + for _ in range(n_in_group): + ra.append(center_ra + rng.uniform(-0.3, 0.3) / 3600) + dec.append(center_dec + rng.uniform(-0.3, 0.3) / 3600) + + # Fill remaining with isolated objects + n_placed = len(ra) + for i in range(n - n_placed): + ra.append(rng.uniform(0, 360)) + dec.append(rng.uniform(-90, 90)) + + diaObjects = _make_diaObjects(ra, dec) + + # Run with small maxSubsetSize (partitioned path) + task_partitioned = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + labels_partitioned = task_partitioned.cluster(diaObjects) + + # Run with large maxSubsetSize (direct path) + task_direct = _make_task(maxClusteringDistance=1.5, maxSubsetSize=50000) + labels_direct = task_direct.cluster(diaObjects) + + # The known tight groups should be clustered in both cases + idx = 0 + for i in range(10): + n_in_group = 0 + group_start = idx + center_ra = float(i * 30) + center_dec = float(i * 10 - 45) + while idx < n_placed: + if (abs(ra[idx] - center_ra) < 1.0 / 3600 + and abs(dec[idx] - center_dec) < 1.0 / 3600): + n_in_group += 1 + idx += 1 + else: + break + + if n_in_group > 1: + # Check direct path clusters them + direct_group = labels_direct.iloc[group_start:group_start + n_in_group] + self.assertEqual(len(direct_group.unique()), 1, + f"Direct path: group {i} should be one cluster") + + # Check partitioned path clusters them + part_group = labels_partitioned.iloc[group_start:group_start + n_in_group] + self.assertEqual(len(part_group.unique()), 1, + f"Partitioned path: group {i} should be one cluster") + + def test_empty_partition_handled(self): + """An empty KMeans partition should be handled gracefully.""" + # We can't easily force an empty partition from KMeans, but we can + # verify the code runs without error on a dataset where some + # partitions might be very small. + rng = np.random.default_rng(7) + # Highly clustered data that may produce uneven partitions + n = 1500 + # Most objects in one location + ra = rng.normal(50.0, 0.01, n).tolist() + dec = rng.normal(20.0, 0.01, n).tolist() + # A few outliers + for i in range(50): + ra.append(float(i * 7)) + dec.append(float(i * 3 - 45)) + + diaObjects = _make_diaObjects(ra + ra[-50:], dec + dec[-50:]) + # Should complete without error + labels = self.task.cluster(diaObjects) + self.assertEqual(len(labels), len(diaObjects)) + + def test_many_partitions(self): + """A dataset requiring many partitions.""" + task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + rng = np.random.default_rng(2024) + n = 5000 # Will create ~5 partitions + + ra = rng.uniform(0, 360, n).tolist() + dec = rng.uniform(-90, 90, n).tolist() + + # Add known duplicates + for i in range(20): + base_ra = rng.uniform(0, 360) + base_dec = rng.uniform(-60, 60) + ra.extend([base_ra, base_ra + 0.2 / 3600]) + dec.extend([base_dec, base_dec + 0.1 / 3600]) + + diaObjects = _make_diaObjects(ra, dec) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels), len(diaObjects)) + + # Verify known duplicates are clustered + for i in range(20): + idx1 = n + 2 * i + idx2 = n + 2 * i + 1 + self.assertEqual(labels.iloc[idx1], labels.iloc[idx2], + f"Known duplicate pair {i} should be clustered") + + +class TestClusterSubset(unittest.TestCase): + """Direct tests of the _cluster_subset helper.""" + + def setUp(self): + self.task = _make_task(maxClusteringDistance=1.5) + + def test_single_object_returns_label_zero(self): + """A single object returns label 0.""" + diaObjects = _make_diaObjects([10.0], [5.0]) + labels = self.task._cluster_subset(diaObjects) + + self.assertEqual(len(labels), 1) + self.assertEqual(labels.iloc[0], 0) + + def test_two_nearby_objects(self): + """Two nearby objects get the same label.""" + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600], + [5.0, 5.0] + ) + labels = self.task._cluster_subset(diaObjects) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_preserves_dataframe_index(self): + """Labels index should match input DataFrame index.""" + diaObjects = _make_diaObjects([10.0, 20.0, 30.0], [5.0, 15.0, 25.0]) + diaObjects.index = pd.Index([5, 10, 15]) + + labels = self.task._cluster_subset(diaObjects) + self.assertTrue(labels.index.equals(diaObjects.index)) + + +class TestClusterEdgeCases(unittest.TestCase): + """Edge cases and special configurations for the cluster method.""" + + def test_all_objects_identical_position(self): + """All objects at the same location should form one cluster.""" + n = 10 + ra = [100.0] * n + dec = [-30.0] * n + diaObjects = _make_diaObjects(ra, dec) + + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), 1) + + def test_negative_declination(self): + """Clustering works for objects in the southern hemisphere.""" + diaObjects = _make_diaObjects( + [200.0, 200.0 + 0.1 / 3600, 200.0 + 0.2 / 3600], + [-60.0, -60.0, -60.0] + ) + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), 1) + + def test_objects_along_equator(self): + """Objects along the celestial equator.""" + # Objects 0.5 arcsec apart along RA at dec=0 + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.5 / 3600, 10.0 + 1.0 / 3600], + [0.0, 0.0, 0.0] + ) + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + # All three within 1 arcsec of each other, should cluster + self.assertEqual(len(labels.unique()), 1) + + def test_minimum_nNeighborsConnectivity(self): + """Task with minimum nNeighborsConnectivity=3 still works.""" + task = _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=3) + + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 20.0, 30.0, 40.0], + [5.0, 5.0, 15.0, 25.0, 35.0] + ) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels), 5) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + def test_maxSubsetSize_boundary(self): + """Dataset exactly at maxSubsetSize should use direct path.""" + n = 1000 + task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + + rng = np.random.default_rng(42) + ra = rng.uniform(0, 360, n).tolist() + dec = rng.uniform(-90, 90, n).tolist() + diaObjects = _make_diaObjects(ra, dec) + + # Should use direct path (n == maxSubsetSize) + labels = task.cluster(diaObjects) + self.assertEqual(len(labels), n) + + def test_maxSubsetSize_plus_one_triggers_partitioning(self): + """Dataset of maxSubsetSize+1 should trigger partitioning.""" + n = 1001 + task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) + + rng = np.random.default_rng(42) + ra = rng.uniform(0, 360, n).tolist() + dec = rng.uniform(-90, 90, n).tolist() + diaObjects = _make_diaObjects(ra, dec) + + labels = task.cluster(diaObjects) + self.assertEqual(len(labels), n) + + def test_three_objects_same_position(self): + """Three objects at the same position form one cluster.""" + diaObjects = _make_diaObjects( + [45.0, 45.0, 45.0], + [-10.0, -10.0, -10.0] + ) + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels.unique()), 1) + + def test_non_default_index(self): + """DataFrame with a non-default integer index.""" + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 50.0], + [5.0, 5.0, 30.0] + ) + diaObjects.index = pd.Index([100, 200, 300]) + + task = _make_task(maxClusteringDistance=1.5) + labels = task.cluster(diaObjects) + + self.assertTrue(labels.index.equals(diaObjects.index)) + self.assertEqual(labels.loc[100], labels.loc[200]) + self.assertNotEqual(labels.loc[100], labels.loc[300]) + + def test_large_nNeighborsConnectivity(self): + """nNeighborsConnectivity larger than dataset is handled.""" + task = _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=30) + + # Only 4 objects but nNeighborsConnectivity=30 + diaObjects = _make_diaObjects( + [10.0, 10.0 + 0.1 / 3600, 50.0, 80.0], + [5.0, 5.0, 30.0, 60.0] + ) + labels = task.cluster(diaObjects) + + self.assertEqual(len(labels), 4) + self.assertEqual(labels.iloc[0], labels.iloc[1]) + + + def setup_module(module): lsst.utils.tests.init() diff --git a/tests/test_deduplicateAllSkyDiaObjects_cluster.py b/tests/test_deduplicateAllSkyDiaObjects_cluster.py deleted file mode 100644 index b8fe450e..00000000 --- a/tests/test_deduplicateAllSkyDiaObjects_cluster.py +++ /dev/null @@ -1,720 +0,0 @@ -# This file is part of ap_association. -# -# Developed for the LSST Data Management System. -# This product includes software developed by the LSST Project -# (https://www.lsst.org). -# See the COPYRIGHT file at the top-level directory of this distribution -# for details of code ownership. -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . - -"""Comprehensive unit tests for the `cluster` method of -DeduplicateAllSkyDiaObjectsTask. -""" - -import numpy as np -import pandas as pd -import tempfile -import os -import unittest - -from lsst.ap.association import DeduplicateAllSkyDiaObjectsTask -from lsst.dax.apdb import ApdbSql -import lsst.utils.tests - - -def _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=30, - maxSubsetSize=50000): - """Create a DeduplicateAllSkyDiaObjectsTask with given config values. - - Creates a temporary SQLite APDB to satisfy the config requirement. - """ - db_file_fd, db_file = tempfile.mkstemp() - apdb_config = ApdbSql.init_database(db_url="sqlite:///" + db_file) - config_file = tempfile.NamedTemporaryFile(delete=False) - apdb_config.save(config_file.name) - config_file.close() - - config = DeduplicateAllSkyDiaObjectsTask.ConfigClass() - config.apdb_config_url = config_file.name - config.maxClusteringDistance = maxClusteringDistance - config.nNeighborsConnectivity = nNeighborsConnectivity - config.maxSubsetSize = maxSubsetSize - - task = DeduplicateAllSkyDiaObjectsTask(config=config) - - os.close(db_file_fd) - os.unlink(db_file) - os.unlink(config_file.name) - - return task - - -def _make_diaObjects(ra_list, dec_list, ids=None): - """Create a minimal DiaObjects DataFrame from lists of ra/dec.""" - n = len(ra_list) - if ids is None: - ids = list(range(1, n + 1)) - df = pd.DataFrame({ - 'diaObjectId': ids, - 'ra': ra_list, - 'dec': dec_list, - }) - df.index = pd.RangeIndex(n) - return df - - -class TestClusterSmallDataset(unittest.TestCase): - """Tests for the cluster method when dataset size <= maxSubsetSize - (direct clustering path via _cluster_subset). - """ - - def setUp(self): - self.task = _make_task(maxClusteringDistance=1.5) - - def test_single_object(self): - """A single object should get a unique cluster label.""" - diaObjects = _make_diaObjects([10.0], [5.0]) - labels = self.task.cluster(diaObjects) - - self.assertEqual(len(labels), 1) - self.assertEqual(labels.iloc[0], 0) - - def test_two_objects_within_threshold(self): - """Two objects within maxClusteringDistance should share a label.""" - # 0.3 arcsec apart in dec - diaObjects = _make_diaObjects( - [10.0, 10.0], - [5.0, 5.0 + 0.3 / 3600] - ) - labels = self.task.cluster(diaObjects) - - self.assertEqual(len(labels), 2) - self.assertEqual(labels.iloc[0], labels.iloc[1]) - - def test_two_objects_beyond_threshold(self): - """Two objects beyond maxClusteringDistance should get different labels.""" - # 10 arcsec apart, well beyond the 1.5 arcsec threshold - diaObjects = _make_diaObjects( - [10.0, 10.0], - [5.0, 5.0 + 10.0 / 3600] - ) - labels = self.task.cluster(diaObjects) - - self.assertEqual(len(labels), 2) - self.assertNotEqual(labels.iloc[0], labels.iloc[1]) - - def test_multiple_distinct_clusters(self): - """Multiple well-separated groups should form distinct clusters.""" - # Three tight groups separated by large distances - ra = [10.0, 10.0 + 0.1 / 3600, # Group 1 - 20.0, 20.0 + 0.2 / 3600, # Group 2 - 30.0, 30.0 + 0.1 / 3600, 30.0 + 0.2 / 3600] # Group 3 - dec = [5.0, 5.0, - 15.0, 15.0, - 25.0, 25.0, 25.0] - - diaObjects = _make_diaObjects(ra, dec) - labels = self.task.cluster(diaObjects) - - self.assertEqual(len(labels), 7) - - # Group 1 should share a label - self.assertEqual(labels.iloc[0], labels.iloc[1]) - # Group 2 should share a label - self.assertEqual(labels.iloc[2], labels.iloc[3]) - # Group 3 should share a label - self.assertEqual(labels.iloc[4], labels.iloc[5]) - self.assertEqual(labels.iloc[5], labels.iloc[6]) - - # Distinct groups should have different labels - unique_group_labels = {labels.iloc[0], labels.iloc[2], labels.iloc[4]} - self.assertEqual(len(unique_group_labels), 3) - - def test_chain_of_objects(self): - """Objects forming a chain: each consecutive pair is within threshold, - but endpoints exceed threshold. - - AgglomerativeClustering with a distance_threshold will merge only - pairs within that distance. With a k-NN connectivity graph, it can - still link objects transitively through neighbors. However, the - distance_threshold limits merging to the configured distance, so - endpoints farther apart than the threshold may end up in separate - clusters depending on the linkage. - """ - # Create a chain of 3 objects, each 0.5 arcsec from its neighbor. - # Total span is 1.0 arcsec < threshold of 1.5 arcsec. - n = 3 - step = 0.5 / 3600 # 0.5 arcsec steps in degrees - ra = [10.0] * n - dec = [5.0 + i * step for i in range(n)] - - diaObjects = _make_diaObjects(ra, dec) - labels = self.task.cluster(diaObjects) - - # All 3 are within 1.0 arcsec of each other, within the 1.5 arcsec - # threshold, so they should form a single cluster. - unique_labels = labels.unique() - self.assertEqual(len(unique_labels), 1, - "Chain of nearby objects within threshold should form a single cluster") - - def test_objects_exactly_at_threshold(self): - """Objects separated by exactly maxClusteringDistance are merged. - - AgglomerativeClustering uses <= for the distance_threshold, so - objects at exactly the threshold distance are merged. - """ - # Exactly 1.5 arcsec apart - diaObjects = _make_diaObjects( - [10.0, 10.0], - [5.0, 5.0 + 1.5 / 3600] - ) - labels = self.task.cluster(diaObjects) - - self.assertEqual(len(labels), 2) - self.assertEqual(labels.iloc[0], labels.iloc[1]) - - def test_returns_series_with_matching_index(self): - """The returned Series should have the same index as the input.""" - diaObjects = _make_diaObjects([10.0, 20.0, 30.0], [5.0, 15.0, 25.0]) - diaObjects.index = pd.Index([10, 20, 30]) - - labels = self.task.cluster(diaObjects) - - self.assertIsInstance(labels, pd.Series) - self.assertTrue(labels.index.equals(diaObjects.index)) - self.assertEqual(labels.name, "cluster_label") - - def test_contiguous_labels(self): - """Cluster labels should be contiguous integers starting from 0.""" - # Create 10 well-separated objects - ra = [float(i * 10) for i in range(10)] - dec = [float(i * 10) for i in range(10)] - diaObjects = _make_diaObjects(ra, dec) - - labels = self.task.cluster(diaObjects) - - unique_labels = sorted(labels.unique()) - self.assertEqual(unique_labels, list(range(len(unique_labels)))) - - def test_large_cluster(self): - """A tight group of many objects should mostly cluster together. - - With k-NN connectivity, very tight groups where all objects are - mutual nearest neighbors should form a single cluster. We use - a small group here to ensure the connectivity graph is dense enough. - """ - rng = np.random.default_rng(42) - n = 15 - # Objects within 0.3 arcsec of center (well within 1.5 arcsec threshold) - center_ra, center_dec = 100.0, -30.0 - ra = center_ra + rng.uniform(-0.3 / 3600, 0.3 / 3600, n) - dec = center_dec + rng.uniform(-0.3 / 3600, 0.3 / 3600, n) - - diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) - labels = self.task.cluster(diaObjects) - - # All should be in the same cluster - self.assertEqual(len(labels.unique()), 1) - - def test_no_duplicates_all_isolated(self): - """When all objects are well separated, each gets its own cluster.""" - n = 20 - # Place objects 1 degree apart - ra = [float(i) for i in range(n)] - dec = [0.0] * n - - diaObjects = _make_diaObjects(ra, dec) - labels = self.task.cluster(diaObjects) - - self.assertEqual(len(labels.unique()), n) - - def test_mixed_duplicates_and_isolated(self): - """A mix of duplicate groups and isolated objects.""" - ra = [10.0, 10.0 + 0.1 / 3600, 10.0 + 0.2 / 3600, # Group of 3 - 50.0, # Isolated - 80.0, 80.0 + 0.1 / 3600, # Pair - 120.0] # Isolated - dec = [5.0, 5.0, 5.0, - 20.0, - 40.0, 40.0, - 60.0] - - diaObjects = _make_diaObjects(ra, dec) - labels = self.task.cluster(diaObjects) - - # Group of 3 should share a label - self.assertEqual(labels.iloc[0], labels.iloc[1]) - self.assertEqual(labels.iloc[1], labels.iloc[2]) - - # Pair should share a label - self.assertEqual(labels.iloc[4], labels.iloc[5]) - - # All distinct clusters - unique_groups = {labels.iloc[0], labels.iloc[3], labels.iloc[4], labels.iloc[6]} - self.assertEqual(len(unique_groups), 4) - - def test_nNeighborsConnectivity_adapted_for_small_n(self): - """When n_objects < nNeighborsConnectivity, the method should adapt.""" - # config nNeighborsConnectivity=30 but only 5 objects - diaObjects = _make_diaObjects( - [10.0, 10.0 + 0.1 / 3600, 20.0, 30.0, 40.0], - [5.0, 5.0, 15.0, 25.0, 35.0] - ) - # Should not raise - labels = self.task.cluster(diaObjects) - self.assertEqual(len(labels), 5) - # First two should be clustered - self.assertEqual(labels.iloc[0], labels.iloc[1]) - - def test_custom_clustering_distance(self): - """Changing maxClusteringDistance affects cluster membership.""" - ra = [10.0, 10.0] - dec = [5.0, 5.0 + 1.0 / 3600] # 1 arcsec apart - - diaObjects = _make_diaObjects(ra, dec) - - # With 1.5 arcsec threshold, should cluster - task_wide = _make_task(maxClusteringDistance=1.5) - labels_wide = task_wide.cluster(diaObjects) - self.assertEqual(labels_wide.iloc[0], labels_wide.iloc[1]) - - # With 0.5 arcsec threshold, should not cluster - task_narrow = _make_task(maxClusteringDistance=0.5) - labels_narrow = task_narrow.cluster(diaObjects) - self.assertNotEqual(labels_narrow.iloc[0], labels_narrow.iloc[1]) - - def test_declination_near_pole(self): - """Objects near the celestial pole should still cluster correctly.""" - # Near north pole - diaObjects = _make_diaObjects( - [0.0, 180.0], # Opposite RA but very close at the pole - [89.9999, 89.9999] - ) - labels = self.task.cluster(diaObjects) - # These are actually close together at the pole - # (cos(dec) factor makes RA separation small) - # but AgglomerativeClustering uses Euclidean on (ra, dec) coords, - # so it will treat them as 180 degrees apart in RA - self.assertEqual(len(labels), 2) - - def test_ra_wraparound(self): - """Objects near RA=0/360 boundary.""" - # Two objects near RA=0, one just above and one just below - diaObjects = _make_diaObjects( - [0.0001, 359.9999], - [5.0, 5.0] - ) - labels = self.task.cluster(diaObjects) - # Euclidean clustering in (ra, dec) space won't handle wraparound; - # these will appear as ~360 degrees apart and won't cluster. - self.assertNotEqual(labels.iloc[0], labels.iloc[1]) - - -class TestClusterLargeDataset(unittest.TestCase): - """Tests for the cluster method when dataset size > maxSubsetSize - (partitioned clustering path via KMeans + boundary merging). - """ - - def setUp(self): - # Use a very small maxSubsetSize to trigger partitioning - self.task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) - - def test_partitioning_triggered(self): - """Verify that partitioning is used for large datasets.""" - rng = np.random.default_rng(42) - n = 2000 - ra = rng.uniform(0, 360, n) - dec = rng.uniform(-90, 90, n) - diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) - - labels = self.task.cluster(diaObjects) - - self.assertEqual(len(labels), n) - # Should have many unique clusters (most objects are isolated) - self.assertGreater(len(labels.unique()), n * 0.9) - - def test_duplicates_found_within_partition(self): - """Duplicates within the same partition should be identified.""" - rng = np.random.default_rng(42) - n = 1500 # Exceeds maxSubsetSize=1000 - - # Place most objects randomly and far apart - ra = rng.uniform(0, 180, n).tolist() - dec = rng.uniform(-45, 45, n).tolist() - - # Add a tight duplicate group (will end up in same partition) - group_center_ra = 90.0 - group_center_dec = 0.0 - for i in range(5): - ra.append(group_center_ra + rng.uniform(-0.3, 0.3) / 3600) - dec.append(group_center_dec + rng.uniform(-0.3, 0.3) / 3600) - - diaObjects = _make_diaObjects(ra, dec) - labels = self.task.cluster(diaObjects) - - # The 5 duplicate objects at the end should share a cluster label - dup_labels = labels.iloc[n:n + 5] - self.assertEqual(len(dup_labels.unique()), 1, - "Tight group within partition should be clustered") - - def test_duplicates_at_partition_boundary(self): - """Duplicates that span a partition boundary should still be merged.""" - # maxSubsetSize=1000, so with 1500 objects we get 2 partitions. - # Create two clusters of objects placed such that KMeans will - # split them into different partitions, with a duplicate pair - # at the boundary. - rng = np.random.default_rng(123) - - # Group A: objects in one region of sky - n_a = 800 - ra_a = rng.uniform(10, 20, n_a).tolist() - dec_a = rng.uniform(10, 20, n_a).tolist() - - # Group B: objects in a different region - n_b = 800 - ra_b = rng.uniform(100, 110, n_b).tolist() - dec_b = rng.uniform(-20, -10, n_b).tolist() - - # Now add a pair of duplicates at the midpoint between groups - # These are close together but KMeans may split them - mid_ra = 55.0 - mid_dec = 0.0 - ra_boundary = [mid_ra, mid_ra + 0.2 / 3600] - dec_boundary = [mid_dec, mid_dec + 0.1 / 3600] - - ra = ra_a + ra_b + ra_boundary - dec = dec_a + dec_b + dec_boundary - - diaObjects = _make_diaObjects(ra, dec) - labels = self.task.cluster(diaObjects) - - # The boundary pair should be clustered together - boundary_labels = labels.iloc[n_a + n_b:] - self.assertEqual(boundary_labels.iloc[0], boundary_labels.iloc[1], - "Boundary duplicate pair should share a cluster label") - - def test_output_labels_are_contiguous(self): - """Labels should be renumbered to be contiguous even after merging.""" - rng = np.random.default_rng(99) - n = 1500 - ra = rng.uniform(0, 360, n) - dec = rng.uniform(-90, 90, n) - diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) - - labels = self.task.cluster(diaObjects) - - unique_labels = sorted(labels.unique()) - expected = list(range(len(unique_labels))) - self.assertEqual(unique_labels, expected) - - def test_returns_series_with_matching_index(self): - """Returned Series matches input DataFrame index.""" - rng = np.random.default_rng(77) - n = 1500 - ra = rng.uniform(0, 360, n) - dec = rng.uniform(-90, 90, n) - diaObjects = _make_diaObjects(ra.tolist(), dec.tolist()) - diaObjects.index = pd.RangeIndex(start=100, stop=100 + n) - - labels = self.task.cluster(diaObjects) - - self.assertIsInstance(labels, pd.Series) - self.assertTrue(labels.index.equals(diaObjects.index)) - self.assertEqual(labels.name, "cluster_label") - - def test_consistency_with_small_path(self): - """Results from the partitioned path should be consistent with - the direct path for datasets that can use either. - """ - rng = np.random.default_rng(55) - n = 1200 - - # Create data with known duplicate structure - # 10 well-separated groups, each with 2-3 tight duplicates - ra = [] - dec = [] - for i in range(10): - center_ra = float(i * 30) - center_dec = float(i * 10 - 45) - n_in_group = rng.integers(2, 4) - for _ in range(n_in_group): - ra.append(center_ra + rng.uniform(-0.3, 0.3) / 3600) - dec.append(center_dec + rng.uniform(-0.3, 0.3) / 3600) - - # Fill remaining with isolated objects - n_placed = len(ra) - for i in range(n - n_placed): - ra.append(rng.uniform(0, 360)) - dec.append(rng.uniform(-90, 90)) - - diaObjects = _make_diaObjects(ra, dec) - - # Run with small maxSubsetSize (partitioned path) - task_partitioned = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) - labels_partitioned = task_partitioned.cluster(diaObjects) - - # Run with large maxSubsetSize (direct path) - task_direct = _make_task(maxClusteringDistance=1.5, maxSubsetSize=50000) - labels_direct = task_direct.cluster(diaObjects) - - # The known tight groups should be clustered in both cases - idx = 0 - for i in range(10): - n_in_group = 0 - group_start = idx - center_ra = float(i * 30) - center_dec = float(i * 10 - 45) - while idx < n_placed: - if (abs(ra[idx] - center_ra) < 1.0 / 3600 - and abs(dec[idx] - center_dec) < 1.0 / 3600): - n_in_group += 1 - idx += 1 - else: - break - - if n_in_group > 1: - # Check direct path clusters them - direct_group = labels_direct.iloc[group_start:group_start + n_in_group] - self.assertEqual(len(direct_group.unique()), 1, - f"Direct path: group {i} should be one cluster") - - # Check partitioned path clusters them - part_group = labels_partitioned.iloc[group_start:group_start + n_in_group] - self.assertEqual(len(part_group.unique()), 1, - f"Partitioned path: group {i} should be one cluster") - - def test_empty_partition_handled(self): - """An empty KMeans partition should be handled gracefully.""" - # We can't easily force an empty partition from KMeans, but we can - # verify the code runs without error on a dataset where some - # partitions might be very small. - rng = np.random.default_rng(7) - # Highly clustered data that may produce uneven partitions - n = 1500 - # Most objects in one location - ra = rng.normal(50.0, 0.01, n).tolist() - dec = rng.normal(20.0, 0.01, n).tolist() - # A few outliers - for i in range(50): - ra.append(float(i * 7)) - dec.append(float(i * 3 - 45)) - - diaObjects = _make_diaObjects(ra + ra[-50:], dec + dec[-50:]) - # Should complete without error - labels = self.task.cluster(diaObjects) - self.assertEqual(len(labels), len(diaObjects)) - - def test_many_partitions(self): - """A dataset requiring many partitions.""" - task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) - rng = np.random.default_rng(2024) - n = 5000 # Will create ~5 partitions - - ra = rng.uniform(0, 360, n).tolist() - dec = rng.uniform(-90, 90, n).tolist() - - # Add known duplicates - for i in range(20): - base_ra = rng.uniform(0, 360) - base_dec = rng.uniform(-60, 60) - ra.extend([base_ra, base_ra + 0.2 / 3600]) - dec.extend([base_dec, base_dec + 0.1 / 3600]) - - diaObjects = _make_diaObjects(ra, dec) - labels = task.cluster(diaObjects) - - self.assertEqual(len(labels), len(diaObjects)) - - # Verify known duplicates are clustered - for i in range(20): - idx1 = n + 2 * i - idx2 = n + 2 * i + 1 - self.assertEqual(labels.iloc[idx1], labels.iloc[idx2], - f"Known duplicate pair {i} should be clustered") - - -class TestClusterSubset(unittest.TestCase): - """Direct tests of the _cluster_subset helper.""" - - def setUp(self): - self.task = _make_task(maxClusteringDistance=1.5) - - def test_single_object_returns_label_zero(self): - """A single object returns label 0.""" - diaObjects = _make_diaObjects([10.0], [5.0]) - labels = self.task._cluster_subset(diaObjects) - - self.assertEqual(len(labels), 1) - self.assertEqual(labels.iloc[0], 0) - - def test_two_nearby_objects(self): - """Two nearby objects get the same label.""" - diaObjects = _make_diaObjects( - [10.0, 10.0 + 0.1 / 3600], - [5.0, 5.0] - ) - labels = self.task._cluster_subset(diaObjects) - self.assertEqual(labels.iloc[0], labels.iloc[1]) - - def test_preserves_dataframe_index(self): - """Labels index should match input DataFrame index.""" - diaObjects = _make_diaObjects([10.0, 20.0, 30.0], [5.0, 15.0, 25.0]) - diaObjects.index = pd.Index([5, 10, 15]) - - labels = self.task._cluster_subset(diaObjects) - self.assertTrue(labels.index.equals(diaObjects.index)) - - -class TestClusterEdgeCases(unittest.TestCase): - """Edge cases and special configurations for the cluster method.""" - - def test_all_objects_identical_position(self): - """All objects at the same location should form one cluster.""" - n = 10 - ra = [100.0] * n - dec = [-30.0] * n - diaObjects = _make_diaObjects(ra, dec) - - task = _make_task(maxClusteringDistance=1.5) - labels = task.cluster(diaObjects) - - self.assertEqual(len(labels.unique()), 1) - - def test_negative_declination(self): - """Clustering works for objects in the southern hemisphere.""" - diaObjects = _make_diaObjects( - [200.0, 200.0 + 0.1 / 3600, 200.0 + 0.2 / 3600], - [-60.0, -60.0, -60.0] - ) - task = _make_task(maxClusteringDistance=1.5) - labels = task.cluster(diaObjects) - - self.assertEqual(len(labels.unique()), 1) - - def test_objects_along_equator(self): - """Objects along the celestial equator.""" - # Objects 0.5 arcsec apart along RA at dec=0 - diaObjects = _make_diaObjects( - [10.0, 10.0 + 0.5 / 3600, 10.0 + 1.0 / 3600], - [0.0, 0.0, 0.0] - ) - task = _make_task(maxClusteringDistance=1.5) - labels = task.cluster(diaObjects) - - # All three within 1 arcsec of each other, should cluster - self.assertEqual(len(labels.unique()), 1) - - def test_minimum_nNeighborsConnectivity(self): - """Task with minimum nNeighborsConnectivity=3 still works.""" - task = _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=3) - - diaObjects = _make_diaObjects( - [10.0, 10.0 + 0.1 / 3600, 20.0, 30.0, 40.0], - [5.0, 5.0, 15.0, 25.0, 35.0] - ) - labels = task.cluster(diaObjects) - - self.assertEqual(len(labels), 5) - self.assertEqual(labels.iloc[0], labels.iloc[1]) - - def test_maxSubsetSize_boundary(self): - """Dataset exactly at maxSubsetSize should use direct path.""" - n = 1000 - task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) - - rng = np.random.default_rng(42) - ra = rng.uniform(0, 360, n).tolist() - dec = rng.uniform(-90, 90, n).tolist() - diaObjects = _make_diaObjects(ra, dec) - - # Should use direct path (n == maxSubsetSize) - labels = task.cluster(diaObjects) - self.assertEqual(len(labels), n) - - def test_maxSubsetSize_plus_one_triggers_partitioning(self): - """Dataset of maxSubsetSize+1 should trigger partitioning.""" - n = 1001 - task = _make_task(maxClusteringDistance=1.5, maxSubsetSize=1000) - - rng = np.random.default_rng(42) - ra = rng.uniform(0, 360, n).tolist() - dec = rng.uniform(-90, 90, n).tolist() - diaObjects = _make_diaObjects(ra, dec) - - labels = task.cluster(diaObjects) - self.assertEqual(len(labels), n) - - def test_three_objects_same_position(self): - """Three objects at the same position form one cluster.""" - diaObjects = _make_diaObjects( - [45.0, 45.0, 45.0], - [-10.0, -10.0, -10.0] - ) - task = _make_task(maxClusteringDistance=1.5) - labels = task.cluster(diaObjects) - - self.assertEqual(len(labels.unique()), 1) - - def test_non_default_index(self): - """DataFrame with a non-default integer index.""" - diaObjects = _make_diaObjects( - [10.0, 10.0 + 0.1 / 3600, 50.0], - [5.0, 5.0, 30.0] - ) - diaObjects.index = pd.Index([100, 200, 300]) - - task = _make_task(maxClusteringDistance=1.5) - labels = task.cluster(diaObjects) - - self.assertTrue(labels.index.equals(diaObjects.index)) - self.assertEqual(labels.loc[100], labels.loc[200]) - self.assertNotEqual(labels.loc[100], labels.loc[300]) - - def test_large_nNeighborsConnectivity(self): - """nNeighborsConnectivity larger than dataset is handled.""" - task = _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=30) - - # Only 4 objects but nNeighborsConnectivity=30 - diaObjects = _make_diaObjects( - [10.0, 10.0 + 0.1 / 3600, 50.0, 80.0], - [5.0, 5.0, 30.0, 60.0] - ) - labels = task.cluster(diaObjects) - - self.assertEqual(len(labels), 4) - self.assertEqual(labels.iloc[0], labels.iloc[1]) - - def test_deterministic_results(self): - """Calling cluster twice on the same data yields identical results.""" - rng = np.random.default_rng(42) - n = 100 - ra = rng.uniform(0, 360, n).tolist() - dec = rng.uniform(-90, 90, n).tolist() - diaObjects = _make_diaObjects(ra, dec) - - task = _make_task(maxClusteringDistance=1.5) - labels1 = task.cluster(diaObjects) - labels2 = task.cluster(diaObjects) - - pd.testing.assert_series_equal(labels1, labels2) - - -def setup_module(module): - lsst.utils.tests.init() - - -if __name__ == "__main__": - lsst.utils.tests.init() - unittest.main() From 30dc8862325846e0914ceef3c727d261f49d8e85 Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Wed, 20 May 2026 14:12:24 -0700 Subject: [PATCH 8/9] wip --- tests/test_deduplicateAllSkyDiaObjects.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_deduplicateAllSkyDiaObjects.py b/tests/test_deduplicateAllSkyDiaObjects.py index 455b7e32..3c949093 100644 --- a/tests/test_deduplicateAllSkyDiaObjects.py +++ b/tests/test_deduplicateAllSkyDiaObjects.py @@ -289,8 +289,6 @@ def test_count_duplicates(self): self.assertEqual(duplicate_count_large, 5) - - def _make_task(maxClusteringDistance=1.5, nNeighborsConnectivity=30, maxSubsetSize=50000): """Create a DeduplicateAllSkyDiaObjectsTask with given config values. @@ -943,7 +941,6 @@ def test_large_nNeighborsConnectivity(self): self.assertEqual(labels.iloc[0], labels.iloc[1]) - def setup_module(module): lsst.utils.tests.init() From 3dbcbfe897a8d5af282b341783f725d1e1b0b30a Mon Sep 17 00:00:00 2001 From: Eric Bellm Date: Wed, 3 Jun 2026 09:37:15 -0700 Subject: [PATCH 9/9] simplify unit tests --- tests/test_deduplicateAllSkyDiaObjects.py | 243 +++++++++------------- 1 file changed, 95 insertions(+), 148 deletions(-) diff --git a/tests/test_deduplicateAllSkyDiaObjects.py b/tests/test_deduplicateAllSkyDiaObjects.py index 3c949093..948cd19f 100644 --- a/tests/test_deduplicateAllSkyDiaObjects.py +++ b/tests/test_deduplicateAllSkyDiaObjects.py @@ -27,7 +27,7 @@ import unittest from lsst.ap.association import DeduplicateAllSkyDiaObjectsTask -from lsst.ap.association.utils import getMidpointFromTimespan, readSchemaFromApdb +from lsst.ap.association.utils import getMidpointFromTimespan from lsst.dax.apdb import Apdb, ApdbSql from lsst.utils import getPackageDir import lsst.utils.tests @@ -54,121 +54,79 @@ def _data_file_name(basename, module_name): class TestDeduplicateAllSkyDiaObjects(unittest.TestCase): + """Tests for the full DeduplicateAllSkyDiaObjectsTask including + run(), remap_clusters(), and count_duplicates(). + """ def setUp(self): - # Create an instance of random generator with fixed seed. + self.task = _make_task(maxClusteringDistance=1.5) + + def _make_apdb_with_duplicates(self): + """Set up an APDB with known duplicate DiaObjects and return + the task, expected duplicate IDs, and APDB access info. + """ rng = np.random.default_rng(1234) - self.db_file_fd, self.db_file = tempfile.mkstemp( - dir=os.path.dirname(__file__)) - self.addCleanup(os.remove, self.db_file) - self.addCleanup(os.close, self.db_file_fd) - - self.apdbConfig = ApdbSql.init_database(db_url="sqlite:///" + self.db_file) - self.config_file = tempfile.NamedTemporaryFile() - self.addCleanup(self.config_file.close) - self.apdbConfig.save(self.config_file.name) - self.apdb = Apdb.from_config(self.apdbConfig) - self.schema = readSchemaFromApdb(self.apdb) - - self.exposure = makeExposure(False, False) - self.regionTime = makeRegionTime(exposure=self.exposure) - self.dateTime = getMidpointFromTimespan(self.regionTime.timespan) - - original_diaObjects = makeDiaObjects(50, self.exposure, rng) - - # create duplicates - dup_0_0 = {"ra": original_diaObjects.iloc[0]['ra'], - "dec": original_diaObjects.iloc[0]['dec'] + 0.2/3600, - "diaObjectId": 1000, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_0_0["%s_psfFluxNdata" % f] = 0 - dup_0_1 = {"ra": original_diaObjects.iloc[0]['ra'] + 0.1/3600, - "dec": original_diaObjects.iloc[0]['dec'] + 0.2/3600, - "diaObjectId": 1001, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_0_1["%s_psfFluxNdata" % f] = 0 - dup_0_2 = {"ra": original_diaObjects.iloc[0]['ra'], - "dec": original_diaObjects.iloc[0]['dec'] - 0.1/3600, - "diaObjectId": 1002, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_0_2["%s_psfFluxNdata" % f] = 0 - dup_0_3 = {"ra": original_diaObjects.iloc[0]['ra'] + 0.2/3600, - "dec": original_diaObjects.iloc[0]['dec'] - 0.1/3600, - "diaObjectId": 1003, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_0_3["%s_psfFluxNdata" % f] = 0 - - dup_1_0 = {"ra": original_diaObjects.iloc[1]['ra'], - "dec": original_diaObjects.iloc[1]['dec'] + 0.2/3600, - "diaObjectId": 2000, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_1_0["%s_psfFluxNdata" % f] = 0 - dup_1_1 = {"ra": original_diaObjects.iloc[1]['ra'] + 0.1/3600, - "dec": original_diaObjects.iloc[1]['dec'] + 0.2/3600, - "diaObjectId": 2001, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_1_1["%s_psfFluxNdata" % f] = 0 - dup_1_2 = {"ra": original_diaObjects.iloc[1]['ra'], - "dec": original_diaObjects.iloc[1]['dec'] - 0.1/3600, - "diaObjectId": 2002, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_1_2["%s_psfFluxNdata" % f] = 0 - - dup_2_0 = {"ra": original_diaObjects.iloc[2]['ra'], - "dec": original_diaObjects.iloc[2]['dec'] + 0.2/3600, - "diaObjectId": 3000, - "nDiaSources": 1} - for f in ["u", "g", "r", "i", "z", "y"]: - dup_2_0["%s_psfFluxNdata" % f] = 0 - - data = [dup_0_0, dup_0_1, dup_0_2, dup_0_3, - dup_1_0, dup_1_1, dup_1_2, - dup_2_0] - duplicate_diaObjects = pd.DataFrame(data=data).set_index("diaObjectId", drop=False) - - self.diaObjects = pd.concat([original_diaObjects, duplicate_diaObjects]) - - self.diaSources = makeDiaSources( - 100, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) - self.diaForcedSources = makeDiaForcedSources( - 200, self.diaObjects["diaObjectId"].to_numpy(), self.exposure, rng) - - # Store the test diaSources as though they were observed a month before - # the current exposure. - dateTime = self.regionTime.timespan.begin.tai - 30 * astropy.units.day - self.apdb.store(dateTime, - self.diaObjects, - self.diaSources, - self.diaForcedSources) - - # These columns are not in the DPDD, yet do appear in DiaSource.yaml. - # We don't need to check them against the default APDB schema. - self.ignoreColumns = ["band", "bboxSize", "isDipole", "flags"] - - def _makeConfig(self, **kwargs): + db_file_fd, db_file = tempfile.mkstemp(dir=os.path.dirname(__file__)) + self.addCleanup(os.remove, db_file) + self.addCleanup(os.close, db_file_fd) + + apdbConfig = ApdbSql.init_database(db_url="sqlite:///" + db_file) + config_file = tempfile.NamedTemporaryFile() + self.addCleanup(config_file.close) + apdbConfig.save(config_file.name) + apdb = Apdb.from_config(apdbConfig) + + exposure = makeExposure(False, False) + regionTime = makeRegionTime(exposure=exposure) + dateTime = getMidpointFromTimespan(regionTime.timespan) + + original_diaObjects = makeDiaObjects(50, exposure, rng) + + duplicate_rows = [] + offsets = [ + (0, [(0, 0.2), (0.1, 0.2), (0, -0.1), (0.2, -0.1)], 1000), + (1, [(0, 0.2), (0.1, 0.2), (0, -0.1)], 2000), + (2, [(0, 0.2)], 3000), + ] + for obj_idx, offset_list, base_id in offsets: + base_ra = original_diaObjects.iloc[obj_idx]['ra'] + base_dec = original_diaObjects.iloc[obj_idx]['dec'] + for i, (dra, ddec) in enumerate(offset_list): + row = {"ra": base_ra + dra / 3600, + "dec": base_dec + ddec / 3600, + "diaObjectId": base_id + i, + "nDiaSources": 1} + for f in ["u", "g", "r", "i", "z", "y"]: + row["%s_psfFluxNdata" % f] = 0 + duplicate_rows.append(row) + + duplicate_diaObjects = pd.DataFrame( + data=duplicate_rows).set_index("diaObjectId", drop=False) + diaObjects = pd.concat([original_diaObjects, duplicate_diaObjects]) + + diaSources = makeDiaSources( + 100, diaObjects["diaObjectId"].to_numpy(), exposure, rng) + diaForcedSources = makeDiaForcedSources( + 200, diaObjects["diaObjectId"].to_numpy(), exposure, rng) + + storeDateTime = regionTime.timespan.begin.tai - 30 * astropy.units.day + apdb.store(storeDateTime, diaObjects, diaSources, diaForcedSources) + config = DeduplicateAllSkyDiaObjectsTask.ConfigClass() - config.apdb_config_url = self.config_file.name - config.earliestMidpointMjdTai = self.dateTime.mjd - config.update(**kwargs) - return config + config.apdb_config_url = config_file.name + config.earliestMidpointMjdTai = dateTime.mjd + task = DeduplicateAllSkyDiaObjectsTask(config=config) + + return task, apdb, regionTime, dateTime def testRun(self): """Test the full run method for deduplication. """ - diaConfig = self._makeConfig() - task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) + task, apdb, regionTime, dateTime = self._make_apdb_with_duplicates() result = task.run() self.assertIsNotNone(result.diaObjectDeduplicationMap) - self.assertEqual(len(result.diaObjectDeduplicationMap), 8) self.assertIn('removedDiaObjectId', result.diaObjectDeduplicationMap.columns) @@ -181,11 +139,9 @@ def testRun(self): kept_ids = set(result.diaObjectDeduplicationMap['keptDiaObjectId']) self.assertTrue(kept_ids.issubset({1, 2, 3})) - # Check that diaSources have been remapped from removed diaObjectIds - # (We can't easily compare to the kept ids because non-duplicated ids - # are not included in the deduplication map.) - updated_diaSources = self.apdb.getDiaSources(region=self.regionTime.region, object_ids=None, - visit_time=self.dateTime, start_time=None) + updated_diaSources = apdb.getDiaSources( + region=regionTime.region, object_ids=None, + visit_time=dateTime, start_time=None) source_object_ids = set(updated_diaSources['diaObjectId']) self.assertTrue(source_object_ids.isdisjoint(removed_ids), "DiaSources should not reference removed diaObjectIds") @@ -193,10 +149,6 @@ def testRun(self): def test_remap_clusters(self): """Test the cluster remapping logic. """ - diaConfig = self._makeConfig() - task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) - - # Create a simple test DataFrame with cluster labels test_data = pd.DataFrame({ 'diaObjectId': [1, 2, 3, 4, 5, 6], 'validityStartMjdTai': [100.0, 101.0, 102.0, 200.0, 201.0, 300.0], @@ -205,11 +157,10 @@ def test_remap_clusters(self): 'dec': [5.0, 5.0, 5.0, 10.0, 10.0, 15.0] }) - dedup_map = task.remap_clusters(test_data) + dedup_map = self.task.remap_clusters(test_data) - self.assertEqual(len(dedup_map), 3) # 2 + 1 remappings + self.assertEqual(len(dedup_map), 3) - # Check that objects 2 and 3 map to object 1 obj2_mapping = dedup_map[dedup_map['removedDiaObjectId'] == 2] obj3_mapping = dedup_map[dedup_map['removedDiaObjectId'] == 3] self.assertEqual(len(obj2_mapping), 1) @@ -217,42 +168,45 @@ def test_remap_clusters(self): self.assertEqual(obj2_mapping.iloc[0]['keptDiaObjectId'], 1) self.assertEqual(obj3_mapping.iloc[0]['keptDiaObjectId'], 1) - # Check that object 5 maps to object 4 obj5_mapping = dedup_map[dedup_map['removedDiaObjectId'] == 5] self.assertEqual(len(obj5_mapping), 1) self.assertEqual(obj5_mapping.iloc[0]['keptDiaObjectId'], 4) def test_cluster(self): - """Test the clustering algorithm. + """Test clustering with known duplicate groups. """ - diaConfig = self._makeConfig(maxClusteringDistance=1.5) - task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) - - cluster_labels = task.cluster(self.diaObjects) - - self.assertEqual(len(cluster_labels), len(self.diaObjects)) + # Three groups of duplicates within 0.3 arcsec, well separated + ra = [10.0, 10.0 + 0.1 / 3600, 10.0 + 0.2 / 3600, + 10.0, 10.0 + 0.2 / 3600, + 50.0, 50.0 + 0.1 / 3600, 50.0 + 0.2 / 3600, + 50.0, + 90.0, 90.0 + 0.2 / 3600] + dec = [5.0, 5.0 + 0.2 / 3600, 5.0 - 0.1 / 3600, + 5.0 + 0.2 / 3600, 5.0 - 0.1 / 3600, + 20.0, 20.0 + 0.2 / 3600, 20.0 - 0.1 / 3600, + 20.0 + 0.2 / 3600, + 40.0, 40.0 + 0.2 / 3600] + diaObjects = _make_diaObjects(ra, dec) + cluster_labels = self.task.cluster(diaObjects) - group1_ids = [1, 1000, 1001, 1002, 1003] - group1_mask = self.diaObjects['diaObjectId'].isin(group1_ids) - group1_labels = cluster_labels[group1_mask] + self.assertEqual(len(cluster_labels), len(diaObjects)) + # Group 1 (indices 0-4) should share a label + group1_labels = cluster_labels.iloc[0:5] self.assertEqual(len(group1_labels.unique()), 1, "Objects in duplicate group 1 should share a cluster label") - group2_ids = [2, 2000, 2001, 2002] - group2_mask = self.diaObjects['diaObjectId'].isin(group2_ids) - group2_labels = cluster_labels[group2_mask] - + # Group 2 (indices 5-8) should share a label + group2_labels = cluster_labels.iloc[5:9] self.assertEqual(len(group2_labels.unique()), 1, "Objects in duplicate group 2 should share a cluster label") - group3_ids = [3, 3000] - group3_mask = self.diaObjects['diaObjectId'].isin(group3_ids) - group3_labels = cluster_labels[group3_mask] - + # Group 3 (indices 9-10) should share a label + group3_labels = cluster_labels.iloc[9:11] self.assertEqual(len(group3_labels.unique()), 1, "Objects in duplicate group 3 should share a cluster label") + # Groups should have different labels self.assertNotEqual(group1_labels.iloc[0], group2_labels.iloc[0], "Group 1 and Group 2 should have different cluster labels") self.assertNotEqual(group1_labels.iloc[0], group3_labels.iloc[0], @@ -263,29 +217,23 @@ def test_cluster(self): def test_count_duplicates(self): """Test the duplicate counting method. """ - diaConfig = self._makeConfig(maxClusteringDistance=1.5) - task = DeduplicateAllSkyDiaObjectsTask(config=diaConfig) - test_catalog = pd.DataFrame({ 'diaObjectId': [1, 2, 3, 4, 5], - 'ra': [10.0, 10.0001, # Pair 1 - 10.1, 10.1001, # Pair 2 - 10.3], # Single + 'ra': [10.0, 10.0001, + 10.1, 10.1001, + 10.3], 'dec': [5.0, 5.0001, 5.1, 5.1001, 5.3] }) - duplicate_count = task.count_duplicates(test_catalog) - + duplicate_count = self.task.count_duplicates(test_catalog) self.assertEqual(duplicate_count, 4) - # Test with custom radius - duplicate_count_small = task.count_duplicates(test_catalog, radius_arcsec=0.1) + duplicate_count_small = self.task.count_duplicates(test_catalog, radius_arcsec=0.1) self.assertEqual(duplicate_count_small, 0) - # Test with large radius - duplicate_count_large = task.count_duplicates(test_catalog, radius_arcsec=3600.0) + duplicate_count_large = self.task.count_duplicates(test_catalog, radius_arcsec=3600.0) self.assertEqual(duplicate_count_large, 5) @@ -451,7 +399,6 @@ def test_returns_series_with_matching_index(self): self.assertTrue(labels.index.equals(diaObjects.index)) self.assertEqual(labels.name, "cluster_label") - def test_large_cluster(self): """A tight group of many objects should mostly cluster together.