From 8196a8ae6ef365407de6e8faea3793dc568492ce Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Tue, 3 Nov 2026 16:22:52 -0800 Subject: [PATCH 1/5] Add MST-based alternative initial grouper Introduces MSTWcsMatch and MSTShearMatch, which replace the footprint/pixelization-based initial clustering with a minimum-spanning-tree equivalent: sources within mst_cut arcseconds (via a chain of neighbours) are placed in the same cluster using scipy connected_components on a sparse near-neighbour graph. This makes initial cluster assignment robust against pixel scale choices while preserving the existing hierarchical confusion-resolution pipeline. New files: - src/hpmcm/mst_utils.py -- mst_group_sources() - src/hpmcm/mst_cell.py -- MSTCellMixin, MSTCellData, MSTShearCellData - src/hpmcm/mst_match.py -- MSTWcsMatch, MSTShearMatch - tests/test_mst_match.py -- 4 tests including a synthetic grouping check Generated with AI Co-Authored-By: SLAC AI --- src/hpmcm/__init__.py | 9 +- src/hpmcm/mst_cell.py | 83 ++++++++++++ src/hpmcm/mst_match.py | 273 ++++++++++++++++++++++++++++++++++++++++ src/hpmcm/mst_utils.py | 58 +++++++++ tests/test_mst_match.py | 173 +++++++++++++++++++++++++ 5 files changed, 595 insertions(+), 1 deletion(-) create mode 100644 src/hpmcm/mst_cell.py create mode 100644 src/hpmcm/mst_match.py create mode 100644 src/hpmcm/mst_utils.py create mode 100644 tests/test_mst_match.py diff --git a/src/hpmcm/__init__.py b/src/hpmcm/__init__.py index a574e27..2bf558a 100644 --- a/src/hpmcm/__init__.py +++ b/src/hpmcm/__init__.py @@ -74,12 +74,14 @@ files. Pass the result directly to ``splitAndReduce``. """ -from . import classify, match_utils, package_utils, shear_utils, utils, viz_utils +from . import classify, match_utils, mst_utils, package_utils, shear_utils, utils, viz_utils from ._version import __version__ from .cell import CellData, ShearCellData from .cluster import ClusterData, ShearClusterData from .footprint import Footprint, FootprintSet from .match import Match +from .mst_cell import MSTCellData, MSTShearCellData +from .mst_match import MSTShearMatch, MSTWcsMatch from .object import ObjectData, ShearObjectData from .output_tables import ( ClusterAssocTable, @@ -114,6 +116,10 @@ "Footprint", "FootprintSet", "Match", + "MSTCellData", + "MSTShearCellData", + "MSTShearMatch", + "MSTWcsMatch", "ObjectAssocTable", "ObjectData", "ObjectShearTable", @@ -129,6 +135,7 @@ "WcsMatch", "classify", "match_utils", + "mst_utils", "package_utils", "shear_utils", "utils", diff --git a/src/hpmcm/mst_cell.py b/src/hpmcm/mst_cell.py new file mode 100644 index 0000000..aa79a19 --- /dev/null +++ b/src/hpmcm/mst_cell.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import numpy as np + +from . import match_utils, mst_utils +from .cell import CellData, ShearCellData +from .footprint import Footprint + + +class MSTCellMixin: + """Mixin that replaces footprint-based ``analyze()`` with MST grouping. + + Mix this before ``CellData`` or ``ShearCellData`` in the MRO to get + MST-based initial clustering. The mixin reads attributes that are + guaranteed to exist on all ``CellData`` subclasses (``n_src``, ``data``, + ``matcher``, ``pixel_match_scale``, ``id_offset``, ``cluster_dict``) and + delegates cluster building to ``_buildClusterData``, which is overridden + by shear subclasses. + """ + + def analyze( # type: ignore[override] + self, + weight_name: str | None = None, + pixel_r2_cut: float = 2.0, + ) -> dict | None: + """Analyze this cell using MST-based source grouping.""" + if self.n_src == 0: # type: ignore[attr-defined] + return None + self._buildMSTClusterData(pixel_r2_cut) + return {} + + def _buildMSTClusterData(self, pixel_r2_cut: float) -> None: + """Group sources via MST and build ClusterData for each component.""" + x_parts: list[np.ndarray] = [] + y_parts: list[np.ndarray] = [] + i_cat_parts: list[np.ndarray] = [] + src_id_parts: list[np.ndarray] = [] + src_idx_parts: list[np.ndarray] = [] + + for i_cat, df in enumerate(self.data): # type: ignore[attr-defined] + x_parts.append(df["x_cell"].values) + y_parts.append(df["y_cell"].values) + i_cat_parts.append(np.full(len(df), i_cat, dtype=np.int32)) + src_id_parts.append(df["id"].values) + src_idx_parts.append(np.arange(len(df), dtype=np.int32)) + + x_all = np.concatenate(x_parts) + y_all = np.concatenate(y_parts) + i_cat_all = np.concatenate(i_cat_parts) + src_id_all = np.concatenate(src_id_parts) + src_idx_all = np.concatenate(src_idx_parts) + + matcher = self.matcher # type: ignore[attr-defined] + cut_pix: float = matcher.mst_cut / matcher.pixToArcsec() + labels = mst_utils.mst_group_sources(x_all, y_all, cut_pix) + + pms: int = self.pixel_match_scale # type: ignore[attr-defined] + id_offset: int = self.id_offset # type: ignore[attr-defined] + + for comp_id in np.unique(labels): + mask = labels == comp_id + sources = np.array( + [i_cat_all[mask], src_id_all[mask], src_idx_all[mask]] + ) + x_comp = x_all[mask] + y_comp = y_all[mask] + x_min_ds = int(np.floor(x_comp.min() / pms)) + x_max_ds = int(np.floor(x_comp.max() / pms)) + 1 + y_min_ds = int(np.floor(y_comp.min() / pms)) + y_max_ds = int(np.floor(y_comp.max() / pms)) + 1 + footprint = Footprint.from_bounds(x_min_ds, x_max_ds, y_min_ds, y_max_ds) + i_cluster = int(comp_id) + id_offset + cluster = self._buildClusterData(i_cluster, footprint, sources) # type: ignore[attr-defined] + self.cluster_dict[i_cluster] = cluster # type: ignore[attr-defined] + match_utils.heirarchicalProcessCluster(cluster, self, pixel_r2_cut) # type: ignore[arg-type] + + +class MSTCellData(MSTCellMixin, CellData): + """CellData that uses MST-based source grouping instead of footprint detection.""" + + +class MSTShearCellData(MSTCellMixin, ShearCellData): + """ShearCellData that uses MST-based source grouping instead of footprint detection.""" diff --git a/src/hpmcm/mst_match.py b/src/hpmcm/mst_match.py new file mode 100644 index 0000000..b7a3b45 --- /dev/null +++ b/src/hpmcm/mst_match.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import dataclasses +import json +from pathlib import Path +from typing import Any + +import numpy as np +import pandas + +from .cell import CellData +from .mst_cell import MSTCellData, MSTShearCellData +from .shear_match import ShearMatch +from .shear_utils import DEFAULT_GEOMETRY, ShearCellGeometry +from .wcs_match import WcsMatch, createGlobalWcs + + +class MSTWcsMatch(WcsMatch): + """WcsMatch that uses MST-based initial grouping instead of footprint detection. + + Sources within ``mst_cut`` arcseconds of each other (via a chain of + neighbours) are placed in the same initial cluster. The subsequent + hierarchical confusion-resolution is identical to the standard + :py:class:`~hpmcm.wcs_match.WcsMatch` pipeline. + + Parameters + ---------- + match_wcs: + Astropy WCS defining the pixel grid. + mst_cut: + Maximum source-to-source separation in **arcseconds** for two sources + to be connected in the initial MST grouping. + **kwargs: + Forwarded to :py:class:`~hpmcm.wcs_match.WcsMatch`. + """ + + def __init__( + self, + match_wcs: Any, + mst_cut: float, + **kwargs: Any, + ) -> None: + self.mst_cut: float = mst_cut + WcsMatch.__init__(self, match_wcs, **kwargs) + + def _buildCellData( + self, + id_offset: int, + corner: np.ndarray, + size: np.ndarray, + idx: int, + ) -> CellData: + return MSTCellData(self, id_offset, corner, size, idx, self.cell_buffer) + + def _geometryDict(self) -> dict[str, Any]: + d = super()._geometryDict() + d["mst_cut"] = float(self.mst_cut) + return d + + @classmethod + def create( + cls, + ref_dir: tuple[float, float], + region_size: tuple[float, float], + pixel_size: float, + mst_cut: float, + **kwargs: Any, + ) -> MSTWcsMatch: + """Create an MSTWcsMatch for a given sky region. + + Parameters + ---------- + ref_dir: + Reference direction (RA, DEC) in degrees. + region_size: + Size of the match region in degrees. + pixel_size: + WCS pixel size in degrees. + mst_cut: + MST grouping radius in arcseconds. + **kwargs: + Forwarded to the constructor. + """ + n_pix = (np.array(region_size) / pixel_size).astype(int) + match_wcs = createGlobalWcs(ref_dir, pixel_size, n_pix) + kwargs.setdefault("n_cell_buffer", 1) + return cls(match_wcs, mst_cut, n_pixels=n_pix, **kwargs) + + @classmethod + def load( + cls, + save_dir: str | Path, + x_range: tuple[int, int] | None = None, + y_range: tuple[int, int] | None = None, + ) -> MSTWcsMatch: + """Restore an MSTWcsMatch from a directory written by save(). + + Parameters + ---------- + save_dir: + Directory written by ``save()``. + x_range: + Optional ``(x_min, x_max)`` inclusive bounds on the cell x-index. + y_range: + Optional ``(y_min, y_max)`` inclusive bounds on the cell y-index. + """ + save_dir = Path(save_dir) + with open(save_dir / "geometry.json") as fh: + geo = json.load(fh) + + ref_dir = tuple(geo["ref_dir"]) + n_pix = np.array(geo["n_pixels"]) + match_wcs = createGlobalWcs( + ref_dir, geo["pixel_size"], n_pix, ctype=geo.get("wcs_ctype", "TAN") + ) + matcher = cls( + match_wcs, + geo["mst_cut"], + n_pixels=n_pix, + cell_size=geo["cell_size"], + cell_buffer=geo["cell_buffer"], + cell_max_object=geo["cell_max_object"], + max_sub_division=geo["max_sub_division"], + pixel_r2_cut=geo["pixel_r2_cut"], + n_cell_buffer=geo["n_cell_buffer"], + pixel_match_scale=geo.get("pixel_match_scale", 1), + ) + matcher.catalog_id_map = {int(k): int(v) for k, v in geo["catalog_id_map"].items()} + + cluster_assoc = pandas.read_parquet(save_dir / "cluster_assoc.parquet") + object_assoc = pandas.read_parquet(save_dir / "object_assoc.parquet") + cluster_stats = pandas.read_parquet(save_dir / "cluster_stats.parquet") + object_stats = pandas.read_parquet(save_dir / "object_stats.parquet") + + cluster_assoc, object_assoc = matcher._filterAssocByRange( + cluster_assoc, object_assoc, x_range, y_range + ) + + matcher._loadReducedData(save_dir) + per_cell_data = matcher._buildPerCellData(cluster_assoc) + matcher._reconstructCells( + cluster_assoc, object_assoc, cluster_stats, object_stats, per_cell_data + ) + return matcher + + +class MSTShearMatch(ShearMatch): + """ShearMatch that uses MST-based initial grouping instead of footprint detection. + + Sources within ``mst_cut`` arcseconds of each other (via a chain of + neighbours) are placed in the same initial cluster. The subsequent + hierarchical confusion-resolution is identical to the standard + :py:class:`~hpmcm.shear_match.ShearMatch` pipeline. + + Parameters + ---------- + mst_cut: + Maximum source-to-source separation in **arcseconds** for two sources + to be connected in the initial MST grouping. + **kwargs: + Forwarded to :py:class:`~hpmcm.shear_match.ShearMatch`. + """ + + def __init__( + self, + mst_cut: float, + **kwargs: Any, + ) -> None: + self.mst_cut: float = mst_cut + ShearMatch.__init__(self, **kwargs) + + def _buildCellData( + self, + id_offset: int, + corner: np.ndarray, + size: np.ndarray, + idx: int, + ) -> CellData: + return MSTShearCellData(self, id_offset, corner, size, idx, self.cell_buffer) + + def _geometryDict(self) -> dict[str, Any]: + d = super()._geometryDict() + d["mst_cut"] = float(self.mst_cut) + return d + + @classmethod + def createShearMatch( # type: ignore[override] + cls, + geometry: ShearCellGeometry = DEFAULT_GEOMETRY, + mst_cut: float = 1.0, + **kwargs: Any, + ) -> MSTShearMatch: + """Create an MSTShearMatch for a given shear cell geometry. + + Parameters + ---------- + geometry: + Cell geometry and WCS matching configuration. + mst_cut: + MST grouping radius in arcseconds. + **kwargs: + Forwarded to the constructor. + """ + kw: dict[str, Any] = dict( + pixel_size=geometry.pixel_size, + n_pixels=geometry.tract_size, + cell_size=geometry.cell_inner_size, + cell_buffer=geometry.match_buffer, + cell_max_object=1000, + n_cell_buffer=1, + geometry=geometry, + ) + kw.update(kwargs) + return cls(mst_cut, **kw) + + @classmethod + def load( # type: ignore[override] + cls, + save_dir: str | Path, + x_range: tuple[int, int] | None = None, + y_range: tuple[int, int] | None = None, + ) -> MSTShearMatch: + """Restore an MSTShearMatch from a directory written by save(). + + Parameters + ---------- + save_dir: + Directory written by ``save()``. + x_range: + Optional ``(x_min, x_max)`` inclusive bounds on the cell x-index. + y_range: + Optional ``(y_min, y_max)`` inclusive bounds on the cell y-index. + """ + save_dir = Path(save_dir) + with open(save_dir / "geometry.json") as fh: + geo = json.load(fh) + + sg = geo["shear_geometry"] + geometry = ShearCellGeometry(**sg) + + matcher = cls( + geo["mst_cut"], + pixel_size=geo["pixel_size"], + n_pixels=np.array(geo["n_pixels"]), + cell_size=geo["cell_size"], + cell_buffer=geo["cell_buffer"], + cell_max_object=geo["cell_max_object"], + max_sub_division=geo["max_sub_division"], + pixel_r2_cut=geo["pixel_r2_cut"], + n_cell_buffer=geo["n_cell_buffer"], + shear_names=geo["shear_names"], + deshear=geo["deshear"], + catalogType=geo["cat_type"], + pixel_match_scale=geo["pixel_match_scale"], + geometry=geometry, + ) + matcher.catalog_id_map = {int(k): int(v) for k, v in geo["catalog_id_map"].items()} + + cluster_assoc = pandas.read_parquet(save_dir / "cluster_assoc.parquet") + object_assoc = pandas.read_parquet(save_dir / "object_assoc.parquet") + cluster_stats = pandas.read_parquet(save_dir / "cluster_stats.parquet") + object_stats = pandas.read_parquet(save_dir / "object_stats.parquet") + + cluster_assoc, object_assoc = matcher._filterAssocByRange( + cluster_assoc, object_assoc, x_range, y_range + ) + + matcher._loadReducedData(save_dir) + per_cell_data = matcher._buildPerCellData(cluster_assoc) + matcher._reconstructCells( + cluster_assoc, object_assoc, cluster_stats, object_stats, per_cell_data + ) + return matcher diff --git a/src/hpmcm/mst_utils.py b/src/hpmcm/mst_utils.py new file mode 100644 index 0000000..0a1dfbd --- /dev/null +++ b/src/hpmcm/mst_utils.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import numpy as np +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import connected_components +from scipy.spatial import cKDTree + + +def mst_group_sources(x: np.ndarray, y: np.ndarray, cut_pix: float) -> np.ndarray: + """Group sources by proximity using MST-equivalent single-linkage clustering. + + Two sources end up in the same group if and only if there is a chain of + sources connecting them where each consecutive pair is within ``cut_pix`` + pixels of each other. This is identical to building the full minimum + spanning tree and cutting every edge longer than ``cut_pix``, but is + computed more efficiently by finding connected components of the sparse + near-neighbour graph. + + Parameters + ---------- + x, y: + Cell-frame pixel coordinates, one entry per source (all catalogs + concatenated). + cut_pix: + Maximum pair distance in cell pixels to keep connected. + + Returns + ------- + Integer component labels with shape ``(N,)``. Sources with the same + label belong to the same group. Labels are 0-indexed and contiguous. + """ + n = len(x) + if n == 0: + return np.array([], dtype=np.int32) + if n == 1: + return np.zeros(1, dtype=np.int32) + + xy = np.column_stack([x, y]) + tree = cKDTree(xy) + pairs = list(tree.query_pairs(r=cut_pix)) + + if not pairs: + return np.arange(n, dtype=np.int32) + + pairs_arr = np.array(pairs) + rows = pairs_arr[:, 0] + cols = pairs_arr[:, 1] + ones = np.ones(len(rows), dtype=np.float32) + + graph = csr_matrix( + ( + np.concatenate([ones, ones]), + (np.concatenate([rows, cols]), np.concatenate([cols, rows])), + ), + shape=(n, n), + ) + _, labels = connected_components(graph, directed=False, return_labels=True) + return labels.astype(np.int32) diff --git a/tests/test_mst_match.py b/tests/test_mst_match.py new file mode 100644 index 0000000..3fe5864 --- /dev/null +++ b/tests/test_mst_match.py @@ -0,0 +1,173 @@ +import glob +import os + +import numpy as np +import pandas + +import hpmcm + +DATADIR = "examples/test_data" +SHEAR_ST = "0p01" +SHEAR = 0.01 +CATALOG_TYPE = "wmom" +TRACT = 10463 + +PIXEL_SIZE = 0.5 / 3600.0 +PIXEL_R2CUT = 4.0 + +REF_DIR = (37.9, 7.0) +REGION_SIZE = (0.375, 0.375) +OUR_GEOM = hpmcm.shear_utils.ShearCellGeometry(ref_dir=REF_DIR) + + +def testMSTWcsMatchBasic(setup_data: int) -> None: + """MSTWcsMatch on a single cell gives reasonable match efficiency.""" + assert setup_data == 0 + + source_tablesfiles = sorted( + glob.glob(os.path.join(DATADIR, f"shear_*_{SHEAR_ST}_cleaned_{TRACT}_ns.pq")) + ) + source_tablesfiles.append(os.path.join(DATADIR, f"object_{TRACT}.pq")) + source_tablesfiles.reverse() + source_tablesfiles = [source_tablesfiles[0], source_tablesfiles[1]] + catalog_ids: list[int] = list(np.arange(len(source_tablesfiles))) + + matcher = hpmcm.MSTWcsMatch.create( + REF_DIR, REGION_SIZE, PIXEL_SIZE, mst_cut=1.0, pixel_r2_cut=PIXEL_R2CUT + ) + matcher.reduceData(source_tablesfiles, catalog_ids) + + assert matcher.n_cell[0] == 4 + assert matcher.n_cell[1] == 4 + + matcher.analysisLoop(range(1, 2), range(1, 2)) + + stats = matcher.extractStats() + assert stats is not None + + obj_lists = hpmcm.classify.classifyObjects(matcher, SNRCut=10.0) + n_good = len(obj_lists["ideal"]) + bad_list = [ + "edge_mixed", "edge_missing", "edge_extra", "orphan", + "missing", "two_missing", "many_missing", "extra", "caught", + ] + n_bad = int(np.sum([len(obj_lists[x]) for x in bad_list])) + assert n_good > 0 + effic = n_good / (n_good + n_bad) + assert effic > 0.75 + + +def testMSTWcsMatchSynthetic() -> None: + """MST groups close sources together and keeps distant sources separate.""" + # 0.5 arcsec/pixel grid; mst_cut=1.5 arcsec -> 3.0 pixels + region_size = (0.1, 0.1) + mst_cut = 1.5 + + matcher = hpmcm.MSTWcsMatch.create( + (0.0, 0.0), region_size, PIXEL_SIZE, mst_cut, pixel_r2_cut=PIXEL_R2CUT + ) + + n_pix = (np.array(region_size) / PIXEL_SIZE).astype(int) + cx, cy = float(n_pix[0]) / 2.0, float(n_pix[1]) / 2.0 + + # Catalogs 0 and 1: sources 0.71 pixels apart — should merge into one cluster + # Catalog 2: source 160+ pixels away — should be its own cluster + df0 = pandas.DataFrame( + {"id": [0], "ra": [0.0], "dec": [0.0], "x_pix": [cx], "y_pix": [cy], "snr": [10.0]} + ) + df1 = pandas.DataFrame( + { + "id": [1], + "ra": [0.0], + "dec": [0.0], + "x_pix": [cx + 0.5], + "y_pix": [cy + 0.5], + "snr": [10.0], + } + ) + df2 = pandas.DataFrame( + { + "id": [2], + "ra": [0.0], + "dec": [0.0], + "x_pix": [cx - 160.0], + "y_pix": [cy - 160.0], + "snr": [10.0], + } + ) + + matcher.setReducedData([df0, df1, df2], [0, 1, 2]) + matcher.analysisLoop() + + stats = matcher.extractStats(central_only=False) + assert stats is not None + + total_objects = sum(len(c.object_dict) for c in matcher.cell_dict.values()) + assert total_objects == 2 + + total_clusters = sum(len(c.cluster_dict) for c in matcher.cell_dict.values()) + assert total_clusters == 2 + + # The object containing catalogs 0 and 1 should have n_src=2 + objects = [ + obj + for cell in matcher.cell_dict.values() + for obj in cell.object_dict.values() + ] + n_src_values = sorted(obj.n_src for obj in objects) + assert n_src_values == [1, 2] + + +def testMSTWcsMatchGeometryDict() -> None: + """mst_cut is persisted in _geometryDict.""" + matcher = hpmcm.MSTWcsMatch.create( + (0.0, 0.0), (0.1, 0.1), PIXEL_SIZE, mst_cut=1.25 + ) + geo = matcher._geometryDict() + assert "mst_cut" in geo + assert geo["mst_cut"] == 1.25 + assert geo["class"] == "MSTWcsMatch" + + +def testMSTShearMatchBasic(setup_data: int) -> None: + """MSTShearMatch on a subset of cells gives reasonable match efficiency.""" + assert setup_data == 0 + + source_tablesfiles = sorted( + glob.glob( + os.path.join(DATADIR, f"shear_{CATALOG_TYPE}_{SHEAR_ST}_uncleaned_{TRACT}_*.pq") + ) + ) + source_tablesfiles.reverse() + catalog_ids: list[int] = list(np.arange(len(source_tablesfiles))) + + matcher = hpmcm.MSTShearMatch.createShearMatch( + geometry=OUR_GEOM, + mst_cut=0.5, + pixel_r2_cut=PIXEL_R2CUT, + deshear=-1 * SHEAR, + ) + matcher.reduceData(source_tablesfiles, catalog_ids) + + assert matcher.n_cell[0] == 201 + assert matcher.n_cell[1] == 201 + + matcher.analysisLoop(range(50, 55), range(170, 175)) + + stats = matcher.extractStats() + shear_stats = matcher.extractShearStats() + + assert stats is not None + assert shear_stats is not None + assert shear_stats["object_shear"] is not None + + obj_lists = hpmcm.classify.classifyObjects(matcher, SNRCut=10.0) + n_good = len(obj_lists["ideal"]) + bad_list = [ + "edge_mixed", "edge_missing", "edge_extra", "orphan", + "missing", "two_missing", "many_missing", "extra", "caught", + ] + n_bad = int(np.sum([len(obj_lists[x]) for x in bad_list])) + assert n_good > 0 + effic = n_good / (n_good + n_bad) + assert effic > 0.80 From 48950c761dcce1551af5c1d88e25c8211676eb80 Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Tue, 3 Nov 2026 16:27:04 -0800 Subject: [PATCH 2/5] =?UTF-8?q?Add=20RunMatchingMST.ipynb=20=E2=80=94=20MS?= =?UTF-8?q?T-matcher=20version=20of=20RunMatching?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop-in replacement for RunMatching.ipynb using MSTShearMatch and MSTWcsMatch. Adds mst_cut_shear and mst_cut_wcs tuneable parameters (in arcseconds) and writes outputs to *_mst_* paths to avoid clobbering standard-matcher results. Generated with AI Co-Authored-By: SLAC AI --- examples/rubin_dp2/RunMatchingMST.ipynb | 367 ++++++++++++++++++++++++ 1 file changed, 367 insertions(+) create mode 100644 examples/rubin_dp2/RunMatchingMST.ipynb diff --git a/examples/rubin_dp2/RunMatchingMST.ipynb b/examples/rubin_dp2/RunMatchingMST.ipynb new file mode 100644 index 0000000..3995a2d --- /dev/null +++ b/examples/rubin_dp2/RunMatchingMST.ipynb @@ -0,0 +1,367 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "60f22d4f-5b19-4316-9636-2eb93d7b8bbb", + "metadata": {}, + "source": [ + "# (re-)Run the matchers — MST version\n", + "\n", + "This is a drop-in replacement for `RunMatching.ipynb` that uses the\n", + "MST-based matchers (`MSTShearMatch`, `MSTWcsMatch`) instead of the\n", + "footprint-based ones.\n", + "\n", + "The key difference is that initial source grouping is done by\n", + "single-linkage clustering at a configurable angular separation\n", + "(`mst_cut`, in arcseconds) rather than by detecting connected pixels\n", + "in a pixelised counts map. This removes the dependence on pixel scale\n", + "choices for the initial grouping step; the subsequent hierarchical\n", + "confusion-resolution is unchanged.\n", + "\n", + "**Tuning parameters added here vs `RunMatching.ipynb`:**\n", + "- `mst_cut_shear` — grouping radius for shear matchers (arcsec)\n", + "- `mst_cut_wcs` — grouping radius for WCS meta-matcher (arcsec)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a074eb6-d9aa-40e1-a442-52262d0fd08b", + "metadata": {}, + "outputs": [], + "source": [ + "# Standard imports\n", + "import tables_io\n", + "import hpmcm\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e03fee38-b519-4708-a051-82b9b6c82db1", + "metadata": {}, + "outputs": [], + "source": [ + "# Set the tract\n", + "tract = 8006\n", + "# This will need to change with tract\n", + "ref_dir=(340.5000000000, -9.6694214876)\n", + "our_geom = hpmcm.shear_utils.ShearCellGeometry(\n", + " ref_dir=ref_dir,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "eddb05fe-f39d-4de4-b040-81a6116c3356", + "metadata": {}, + "outputs": [], + "source": [ + "# If this is True only run on a single cell\n", + "test_only = False" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890", + "metadata": {}, + "outputs": [], + "source": [ + "# MST grouping radii (arcseconds)\n", + "# mst_cut_shear: for shear catalogs (0.2 arcsec/pixel native scale)\n", + "# comparable to pixel_match_scale=2 footprint grouping (~0.4-0.8 arcsec)\n", + "mst_cut_shear = 0.8\n", + "\n", + "# mst_cut_wcs: for the WCS meta-matcher (0.2 arcsec/pixel, pixel_match_scale=4)\n", + "# comparable to pixel_match_scale=4 footprint grouping (~0.8-1.6 arcsec)\n", + "mst_cut_wcs = 1.5" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3edac63a-e803-48a0-8ad1-cb09c5e76b27", + "metadata": {}, + "outputs": [], + "source": [ + "# Files and catalogs ids for rubin metadetect\n", + "rubin_shear_match_files = [\n", + " f'{tract}/shear_object_reduced_uncleaned_{tract}_ns.parq',\n", + " f'{tract}/shear_object_reduced_uncleaned_{tract}_1m.parq',\n", + " f'{tract}/shear_object_reduced_uncleaned_{tract}_1p.parq',\n", + " f'{tract}/shear_object_reduced_uncleaned_{tract}_2m.parq',\n", + " f'{tract}/shear_object_reduced_uncleaned_{tract}_2p.parq', \n", + "]\n", + "rubin_shear_match_ids = list(range(len(rubin_shear_match_files)))\n", + "\n", + "# Files and catalogs ids for desc metadetect\n", + "desc_shear_match_files = [\n", + " f'{tract}/mdet-run-reduced_uncleaned_{tract}_ns.parq',\n", + " f'{tract}/mdet-run-reduced_uncleaned_{tract}_1m.parq',\n", + " f'{tract}/mdet-run-reduced_uncleaned_{tract}_1p.parq',\n", + "]\n", + "desc_shear_match_ids = list(range(len(desc_shear_match_files)))\n", + "\n", + "# Files and catalogs ids for meta-matching\n", + "all_files = [\n", + " f'{tract}/object_reduced_{tract}.parq',\n", + " f'{tract}/anacal_reduced_{tract}.parq',\n", + " f'{tract}/desc_mst_md_match_object_joined.pq',\n", + " f'{tract}/rubin_mst_md_match_object_joined.pq',\n", + "]\n", + "all_ids = list(range(len(all_files)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ffd22736-c485-4bf7-acc9-6c0736513b7d", + "metadata": {}, + "outputs": [], + "source": [ + "# Allow for re-running various matchers\n", + "rematch_rubin_shear = True\n", + "rematch_desc_shear = True\n", + "rematch_objects = True\n", + "\n", + "if rematch_rubin_shear:\n", + " rubin_shear_match = hpmcm.MSTShearMatch.createShearMatch(\n", + " geometry=our_geom,\n", + " mst_cut=mst_cut_shear,\n", + " pixel_match_scale=2,\n", + " pixel_r2_cut=16,\n", + " deshear=-0.01,\n", + " shear_names=['ns', '1m', '1p', '2m', '2p'], \n", + " )\n", + " rubin_shear_match.reduceData(rubin_shear_match_files, rubin_shear_match_ids)\n", + " if test_only:\n", + " rubin_shear_match.analysisLoop(range(100, 101), range(100, 101))\n", + " else:\n", + " rubin_shear_match.analysisLoop()\n", + " rubin_stats = rubin_shear_match.extractStats()\n", + " rubin_stats.update(**rubin_shear_match.extractShearStats())\n", + " tables_io.write(rubin_stats, f'{tract}/rubin_mst_md_match_')\n", + " rubin_shear_match.save(f'{tract}/rubin_mst_md_match')\n", + " \n", + "\n", + "if rematch_desc_shear:\n", + " desc_shear_match = hpmcm.MSTShearMatch.createShearMatch(\n", + " geometry=our_geom,\n", + " mst_cut=mst_cut_shear,\n", + " deshear=-0.01,\n", + " pixel_match_scale=2,\n", + " pixel_r2_cut=16, \n", + " shear_names=['ns', '1m', '1p'],\n", + " )\n", + " desc_shear_match.reduceData(desc_shear_match_files, desc_shear_match_ids)\n", + " if test_only: \n", + " desc_shear_match.analysisLoop(range(100, 101), range(100, 101))\n", + " else:\n", + " desc_shear_match.analysisLoop()\n", + " desc_stats = desc_shear_match.extractStats()\n", + " desc_stats.update(**desc_shear_match.extractShearStats())\n", + " tables_io.write(desc_stats, f'{tract}/desc_mst_md_match_') \n", + " desc_shear_match.save(f'{tract}/desc_mst_md_match')\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "608b0387-feb2-42e1-aa47-bf93ceb39030", + "metadata": {}, + "outputs": [], + "source": [ + "shear_rubin_obj_stats = tables_io.read(f'{tract}/rubin_mst_md_match_object_stats.pq')\n", + "shear_rubin_obj_assoc = tables_io.read(f'{tract}/rubin_mst_md_match_object_assoc.pq')\n", + "shear_rubin_obj_shear = tables_io.read(f'{tract}/rubin_mst_md_match_object_shear.pq')\n", + "\n", + "joined_rubin_shear_obj = hpmcm.output_tables.buildJoinedObjectTable(\n", + " shear_rubin_obj_stats,\n", + " shear_rubin_obj_assoc,\n", + " rubin_shear_match_files,\n", + " rubin_shear_match_ids,\n", + " object_shear=f'{tract}/rubin_mst_md_match_object_shear.pq',\n", + ")\n", + "stats_cols = [f\"{band}_gaussFlux\" for band in 'griz']\n", + "stats_cols += [f\"{band}_gaussFluxErr\" for band in 'griz']\n", + "drop_cols = []\n", + "for col_prefix in [\n", + " 'tract', 'ra', 'dec', 'gauss_snr', 'patch_x', 'patch_y', \n", + " 'cell_idx_x', 'cell_idx_y', 'shear', 'meta_step',\n", + "]:\n", + " drop_cols += [f\"{col_prefix}_{id}\" for id in rubin_shear_match_ids]\n", + "\n", + "rd_joined_rubin_shear_obj = hpmcm.reduceJoinedTable(\n", + " joined_rubin_shear_obj,\n", + " catalog_ids=rubin_shear_match_ids,\n", + " drop_cols=drop_cols,\n", + " stats_cols=stats_cols,\n", + ")\n", + "tables_io.write(rd_joined_rubin_shear_obj, f'{tract}/rubin_mst_md_match_object_joined.pq')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0ccc0274-d134-434b-9865-31e2a9416003", + "metadata": {}, + "outputs": [], + "source": [ + "shear_desc_obj_stats = tables_io.read(f'{tract}/desc_mst_md_match_object_stats.pq')\n", + "shear_desc_obj_assoc = tables_io.read(f'{tract}/desc_mst_md_match_object_assoc.pq')\n", + "shear_desc_obj_shear = tables_io.read(f'{tract}/desc_mst_md_match_object_shear.pq')\n", + "\n", + "joined_desc_shear_obj = hpmcm.output_tables.buildJoinedObjectTable(\n", + " shear_desc_obj_stats,\n", + " shear_desc_obj_assoc,\n", + " desc_shear_match_files,\n", + " desc_shear_match_ids,\n", + " object_shear=f'{tract}/desc_mst_md_match_object_shear.pq',\n", + ")\n", + "stats_cols = [f\"flux_{band}\" for band in 'riz']\n", + "stats_cols += [f\"flux_err_{band}\" for band in 'riz']\n", + "drop_cols = []\n", + "for col_prefix in [\n", + " 'ra', 'dec', 'gauss_snr', 'patch_x', 'patch_y', \n", + " 'cell_idx_x', 'cell_idx_y', 'shear', 'meta_step',\n", + "]:\n", + " drop_cols += [f\"{col_prefix}_{id}\" for id in rubin_shear_match_ids]\n", + "\n", + "rd_joined_desc_shear_obj = hpmcm.reduceJoinedTable(\n", + " joined_desc_shear_obj,\n", + " catalog_ids=desc_shear_match_ids,\n", + " drop_cols=drop_cols,\n", + " stats_cols=stats_cols,\n", + ")\n", + "tables_io.write(rd_joined_desc_shear_obj, f'{tract}/desc_mst_md_match_object_joined.pq')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "408c4f7b-3b36-46ac-b744-b5340c2e60c1", + "metadata": {}, + "outputs": [], + "source": [ + "# Draw a cluster from the rubin shear matching\n", + "# countsMap() re-computes the pixel map from source positions for visualisation\n", + "if rematch_rubin_shear:\n", + " cell_data = rubin_shear_match.cell_dict[20200]\n", + " cluster_id = list(cell_data.cluster_dict.keys())[1]\n", + " cluster_data = cell_data.cluster_dict[cluster_id]\n", + " img = cell_data.countsMap()\n", + " _ = hpmcm.viz_utils.showCluster(\n", + " img,\n", + " cluster_data,\n", + " cell_data,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3328791e-876d-4a4d-8982-4f04afaaab29", + "metadata": {}, + "outputs": [], + "source": [ + "# Draw a cluster from the desc shear matching\n", + "if rematch_desc_shear:\n", + " cell_data = desc_shear_match.cell_dict[20200]\n", + " cluster_id = list(cell_data.cluster_dict.keys())[1]\n", + " cluster_data = cell_data.cluster_dict[cluster_id]\n", + " img = cell_data.countsMap()\n", + " _ = hpmcm.viz_utils.showCluster(\n", + " img,\n", + " cluster_data,\n", + " cell_data,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5bd02341-0983-4dc5-85ab-368230a0ac41", + "metadata": {}, + "outputs": [], + "source": [ + "if rematch_objects:\n", + " pixel_scale = 0.2\n", + " n_pix_region = 30000\n", + " cell_size = 150\n", + " region_size_deg = n_pix_region*pixel_scale/3600\n", + " mst_wcs_match = hpmcm.MSTWcsMatch.create(\n", + " ref_dir=ref_dir,\n", + " region_size=(region_size_deg, region_size_deg),\n", + " pixel_size=pixel_scale/3600,\n", + " mst_cut=mst_cut_wcs,\n", + " cell_size=cell_size,\n", + " cell_buffer=5,\n", + " cell_max_object=10000,\n", + " pixel_match_scale=4,\n", + " pixel_r2_cut=16, \n", + " )\n", + " \n", + " mst_wcs_match.reduceData(all_files, all_ids)\n", + " if test_only:\n", + " mst_wcs_match.analysisLoop(range(100, 101), range(100, 101))\n", + " else:\n", + " mst_wcs_match.analysisLoop()\n", + " stats = mst_wcs_match.extractStats()\n", + " tables_io.write(stats, f'{tract}/all_mst_match_') \n", + " mst_wcs_match.save(f'{tract}/all_mst_match')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c900bb1b-8605-4cf2-92bf-5d1991a925d8", + "metadata": {}, + "outputs": [], + "source": [ + "# Draw a cluster from the meta matching\n", + "if rematch_objects:\n", + " cell_data = mst_wcs_match.cell_dict[20200]\n", + " cluster_id = list(cell_data.cluster_dict.keys())[0]\n", + " cluster_data = cell_data.cluster_dict[cluster_id]\n", + " img = cell_data.countsMap()\n", + " _ = hpmcm.viz_utils.showCluster(\n", + " img,\n", + " cluster_data,\n", + " cell_data,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ee7ec80d-7681-44b5-92bf-5f5a4a82e6fa", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbformat_minor": 5, + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From f77c6a1fa495b49459570833e5e7f5f4aa1edf21 Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Tue, 3 Nov 2026 16:49:39 -0800 Subject: [PATCH 3/5] Replace pixel-zoom splitting with MST-native subdivision mstProcessCluster, mstProcessObject, and mstSplitObject replace the heirarchical* pixel-zoom functions for MST cells. Splitting now re-runs mst_group_sources with mst_cut halved at each recursion level, keeping the entire pipeline free of pixelization choices. Generated with AI Co-Authored-By: SLAC AI --- examples/rubin_dp2/RunMatchingMST.ipynb | 35 +++- src/hpmcm/mst_cell.py | 4 +- src/hpmcm/mst_utils.py | 210 ++++++++++++++++++++++++ 3 files changed, 240 insertions(+), 9 deletions(-) diff --git a/examples/rubin_dp2/RunMatchingMST.ipynb b/examples/rubin_dp2/RunMatchingMST.ipynb index 3995a2d..a114046 100644 --- a/examples/rubin_dp2/RunMatchingMST.ipynb +++ b/examples/rubin_dp2/RunMatchingMST.ipynb @@ -25,7 +25,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "5a074eb6-d9aa-40e1-a442-52262d0fd08b", "metadata": {}, "outputs": [], @@ -39,7 +39,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "e03fee38-b519-4708-a051-82b9b6c82db1", "metadata": {}, "outputs": [], @@ -55,7 +55,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "eddb05fe-f39d-4de4-b040-81a6116c3356", "metadata": {}, "outputs": [], @@ -66,7 +66,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890", "metadata": {}, "outputs": [], @@ -83,7 +83,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "3edac63a-e803-48a0-8ad1-cb09c5e76b27", "metadata": {}, "outputs": [], @@ -121,7 +121,19 @@ "execution_count": null, "id": "ffd22736-c485-4bf7-acc9-6c0736513b7d", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "......... 10!\n", + "......... 20!\n", + "......... 30!\n", + "......... 40!\n", + ".." + ] + } + ], "source": [ "# Allow for re-running various matchers\n", "rematch_rubin_shear = True\n", @@ -342,6 +354,14 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c984950d-375a-4de5-a1c7-0affd74cc6f7", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -358,7 +378,8 @@ "file_extension": ".py", "mimetype": "text/x-python", "name": "python", - "nbformat_minor": 5, + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", "version": "3.12.12" } }, diff --git a/src/hpmcm/mst_cell.py b/src/hpmcm/mst_cell.py index aa79a19..ace62f4 100644 --- a/src/hpmcm/mst_cell.py +++ b/src/hpmcm/mst_cell.py @@ -2,7 +2,7 @@ import numpy as np -from . import match_utils, mst_utils +from . import mst_utils from .cell import CellData, ShearCellData from .footprint import Footprint @@ -72,7 +72,7 @@ def _buildMSTClusterData(self, pixel_r2_cut: float) -> None: i_cluster = int(comp_id) + id_offset cluster = self._buildClusterData(i_cluster, footprint, sources) # type: ignore[attr-defined] self.cluster_dict[i_cluster] = cluster # type: ignore[attr-defined] - match_utils.heirarchicalProcessCluster(cluster, self, pixel_r2_cut) # type: ignore[arg-type] + mst_utils.mstProcessCluster(cluster, self, pixel_r2_cut, cut_pix) # type: ignore[arg-type] class MSTCellData(MSTCellMixin, CellData): diff --git a/src/hpmcm/mst_utils.py b/src/hpmcm/mst_utils.py index 0a1dfbd..dc60425 100644 --- a/src/hpmcm/mst_utils.py +++ b/src/hpmcm/mst_utils.py @@ -1,10 +1,17 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np from scipy.sparse import csr_matrix from scipy.sparse.csgraph import connected_components from scipy.spatial import cKDTree +if TYPE_CHECKING: + from .cell import CellData + from .cluster import ClusterData + from .object import ObjectData + def mst_group_sources(x: np.ndarray, y: np.ndarray, cut_pix: float) -> np.ndarray: """Group sources by proximity using MST-equivalent single-linkage clustering. @@ -56,3 +63,206 @@ def mst_group_sources(x: np.ndarray, y: np.ndarray, cut_pix: float) -> np.ndarra ) _, labels = connected_components(graph, directed=False, return_labels=True) return labels.astype(np.int32) + + +def mstProcessCluster( + cluster: ClusterData, + cell_data: CellData, + pixel_r2_cut: float, + mst_cut_pix: float, +) -> list[ObjectData]: + """Process a cluster into objects using MST-based splitting. + + Replaces :py:func:`~hpmcm.match_utils.heirarchicalProcessCluster` for + MST matchers. Computes cluster centroid statistics then delegates + per-object refinement to :py:func:`mstProcessObject`. + + Parameters + ---------- + cluster: + Cluster to process. + cell_data: + Parent cell. + pixel_r2_cut: + Maximum squared distance from object centroid (cell pixels) for a + source to be considered well-matched. + mst_cut_pix: + Initial MST grouping radius in cell pixels. Halved at each + recursion level by :py:func:`mstSplitObject`. + """ + cluster.extract(cell_data) + assert cluster.data is not None + + cluster.n_src = len(cluster.data.i_cat) + cluster.n_unique = len(np.unique(cluster.data.i_cat.values)) + + if cluster.n_src == 0: + print("Empty cluster", cluster.n_src, cluster.n_unique) + return cluster.objects + + if cluster.n_src == 1: + cluster.x_cent = cluster.data.x_cell.values[0] + cluster.y_cent = cluster.data.y_cell.values[0] + cluster.dist_2 = np.zeros((1)) + cluster.rms_dist = 0.0 + cluster.snr_mean = cluster.data.snr.values[0] + cluster.snr_rms = 0.0 + initial_object = cluster.addObject(cell_data) + mstProcessObject(initial_object, cell_data, pixel_r2_cut, mst_cut_pix) + return cluster.objects + + sum_snr = np.sum(cluster.data.snr) + cluster.x_cent = np.sum(cluster.data.x_cell * cluster.data.snr) / sum_snr + cluster.y_cent = np.sum(cluster.data.y_cell * cluster.data.snr) / sum_snr + cluster.dist_2 = ( + (cluster.x_cent - cluster.data.x_cell) ** 2 + + (cluster.y_cent - cluster.data.y_cell) ** 2 + ) + cluster.rms_dist = np.sqrt(np.mean(cluster.dist_2)) + cluster.snr_mean = np.mean(cluster.data.snr.values) + cluster.snr_rms = np.std(cluster.data.snr.values) + + initial_object = cluster.addObject(cell_data) + mstProcessObject(initial_object, cell_data, pixel_r2_cut, mst_cut_pix) + return cluster.objects + + +def mstProcessObject( + obj_data: ObjectData, + cell_data: CellData, + pixel_r2_cut: float, + mst_cut_pix: float, + recurse: int = 0, +) -> None: + """Refine one object, splitting it if sources are confused or out of range. + + Replaces :py:func:`~hpmcm.match_utils.heirarchicalProcessObject` for + MST matchers. When splitting is needed, delegates to + :py:func:`mstSplitObject` with ``mst_cut_pix / 2``. + + Parameters + ---------- + obj_data: + Object to refine. + cell_data: + Parent cell. + pixel_r2_cut: + Maximum squared distance from centroid (cell pixels) for a source + to be well-matched. + mst_cut_pix: + MST cut radius in cell pixels for the *current* recursion level. + :py:func:`mstSplitObject` will be called with ``mst_cut_pix / 2``. + recurse: + Current recursion depth. + """ + if recurse > cell_data.matcher.max_sub_division: + return + obj_data.recurse = recurse + + if obj_data.n_src == 0: + print("Empty object", obj_data.n_src, obj_data.n_unique, recurse) + return + + assert obj_data.data is not None + if obj_data.mask.sum() == 1: + obj_data.x_cent = obj_data.data.x_cell.values[0] + obj_data.y_cent = obj_data.data.y_cell.values[0] + obj_data.dist_2 = np.zeros((1), float) + obj_data.rms_dist = 0.0 + obj_data.snr_mean = obj_data.data.snr.values[0] + obj_data.snr_rms = 0.0 + return + + sum_snr = np.sum(obj_data.data.snr) + obj_data.x_cent = np.sum(obj_data.data.x_cell * obj_data.data.snr) / sum_snr + obj_data.y_cent = np.sum(obj_data.data.y_cell * obj_data.data.snr) / sum_snr + obj_data.dist_2 = np.array( + (obj_data.x_cent - obj_data.data.x_cell) ** 2 + + (obj_data.y_cent - obj_data.data.y_cell) ** 2 + ) + obj_data.rms_dist = np.sqrt(np.mean(obj_data.dist_2)) + obj_data.snr_mean = np.mean(obj_data.data.snr.values) + obj_data.snr_rms = np.std(obj_data.data.snr.values) + + sub_mask = obj_data.dist_2 < pixel_r2_cut + if (obj_data.n_unique == obj_data.n_src) and sub_mask.all(): + return + + if recurse >= cell_data.matcher.max_sub_division: + return + + mstSplitObject(obj_data, cell_data, pixel_r2_cut, mst_cut_pix / 2, recurse=recurse + 1) + + +def mstSplitObject( + obj_data: ObjectData, + cell_data: CellData, + pixel_r2_cut: float, + mst_cut_pix: float, + recurse: int = 0, +) -> None: + """Split a confused object by re-grouping its sources at a finer MST cut. + + Replaces :py:func:`~hpmcm.match_utils.heirarchicalSplitObject` for + MST matchers. Runs :py:func:`mst_group_sources` with ``mst_cut_pix`` + on the sources in the object. If the sources fall into more than one + group the object is divided into sub-objects, each of which is then + refined by :py:func:`mstProcessObject` at the same ``mst_cut_pix``. + If only one group is found, the cut is halved and the split is retried. + + Parameters + ---------- + obj_data: + Object to split. + cell_data: + Parent cell. + pixel_r2_cut: + Maximum squared distance from centroid (cell pixels). + mst_cut_pix: + MST grouping radius in cell pixels for this split attempt. + Halved on each retry. + recurse: + Current recursion depth. + """ + if recurse > cell_data.matcher.max_sub_division: + return + obj_data.recurse = recurse + obj_data.extract() + assert obj_data.data is not None + + x = obj_data.data.x_cell.values + y = obj_data.data.y_cell.values + labels = mst_group_sources(x, y, mst_cut_pix) + n_groups = len(np.unique(labels)) + + if n_groups == 1: + if recurse >= cell_data.matcher.max_sub_division: + return + mstSplitObject(obj_data, cell_data, pixel_r2_cut, mst_cut_pix / 2, recurse=recurse + 1) + return + + biggest_label = int(np.argmax(np.bincount(labels))) + biggest_mask = np.zeros(obj_data.mask.shape, dtype=bool) + + for label in np.unique(labels): + sub_mask_local = labels == label + + # Map from indices-within-object back to indices-within-cluster + new_mask = np.zeros(obj_data.mask.shape, dtype=bool) + count = 0 + for i, val in enumerate(obj_data.mask): + if val: + new_mask[i] = sub_mask_local[count] + count += 1 + assert sub_mask_local.sum() == new_mask.sum() + + if label == biggest_label: + biggest_mask = new_mask + continue + + new_object = obj_data.parent_cluster.addObject(cell_data, new_mask) + mstProcessObject(new_object, cell_data, pixel_r2_cut, mst_cut_pix, recurse=recurse) + + obj_data.mask = biggest_mask + obj_data.extract() + mstProcessObject(obj_data, cell_data, pixel_r2_cut, mst_cut_pix, recurse=recurse) From 4f61e070abcfd84160211a1b9428f55fa046d996 Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Tue, 3 Nov 2026 16:56:44 -0800 Subject: [PATCH 4/5] Add showMSTCluster and showMSTObjects visualization functions Draw MST links between sources instead of a pixel counts-map background. showMSTCluster colours nodes by catalog ID; showMSTObjects colours intra-object edges solid by object and inter-object (cut) edges dashed grey. Generated with AI Co-Authored-By: SLAC AI --- src/hpmcm/viz_utils.py | 123 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/src/hpmcm/viz_utils.py b/src/hpmcm/viz_utils.py index 830e63d..b25d2df 100644 --- a/src/hpmcm/viz_utils.py +++ b/src/hpmcm/viz_utils.py @@ -3,11 +3,27 @@ import matplotlib.pyplot as plt import numpy as np from matplotlib.figure import Figure, SubFigure +from scipy.sparse import csr_matrix +from scipy.sparse.csgraph import minimum_spanning_tree +from scipy.spatial.distance import pdist, squareform from .cell import CellData from .cluster import ClusterData from .match import Match +_CATALOG_COLORS = ["C0", "C1", "C2", "C3", "C4"] +_OBJECT_COLORS = ["tab:red", "tab:blue", "tab:green", "tab:cyan", "tab:orange"] + + +def _mst_edges(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Return (row, col) index arrays of the MST edges for the given positions.""" + if len(x) < 2: + return np.empty(0, dtype=int), np.empty(0, dtype=int) + dist_mat = csr_matrix(squareform(pdist(np.column_stack([x, y])))) + mst = minimum_spanning_tree(dist_mat) + rows, cols = mst.nonzero() + return rows.astype(int), cols.astype(int) + def showShearObjs(matcher: Match, i_k: tuple[int, int]) -> Figure | SubFigure: """Draw the objects in a cluster @@ -254,3 +270,110 @@ def showObjectsV2( img.axes.set_ylabel("y [cell pixels]") assert img.axes.figure is not None return img.axes.figure + + +def showMSTCluster( + cluster: ClusterData, + cell_data: CellData, +) -> Figure | SubFigure: + """Draw a cluster using MST links instead of a pixel image. + + Sources are coloured by catalog ID. MST edges connect adjacent + sources in the minimum spanning tree of the cluster. The + SNR-weighted cluster centroid is marked with a black cross. + + Parameters + ---------- + cluster: + Cluster to visualise. + cell_data: + Parent cell. + + Returns + ------- + Figure containing the plot. + """ + cluster.extract(cell_data) + assert cluster.data is not None + + x = cluster.data.x_cell.values + y = cluster.data.y_cell.values + i_cat = cluster.data.i_cat.values + + fig, ax = plt.subplots() + + rows, cols = _mst_edges(x, y) + for r, c in zip(rows, cols): + ax.plot([x[r], x[c]], [y[r], y[c]], color="grey", linewidth=0.8, zorder=1) + + for cat_id in np.unique(i_cat): + m = i_cat == cat_id + ax.scatter( + x[m], + y[m], + c=_CATALOG_COLORS[int(cat_id) % len(_CATALOG_COLORS)], + label=f"catalog {int(cat_id)}", + zorder=2, + ) + + ax.scatter(cluster.x_cent, cluster.y_cent, marker="+", c="black", s=100, zorder=3) + ax.set_xlabel("x [cell pixels]") + ax.set_ylabel("y [cell pixels]") + ax.legend(loc="best", fontsize="small") + return fig + + +def showMSTObjects( + cluster: ClusterData, + cell_data: CellData, +) -> Figure | SubFigure: + """Draw a cluster's objects using MST links instead of a pixel image. + + Intra-object MST edges are drawn as solid lines coloured by object. + Inter-object edges (cut by the MST splitting step) are drawn as + dashed grey lines. Object centroids are marked with crosses. + + Parameters + ---------- + cluster: + Cluster to visualise. + cell_data: + Parent cell. + + Returns + ------- + Figure containing the plot. + """ + cluster.extract(cell_data) + assert cluster.data is not None + + x = cluster.data.x_cell.values + y = cluster.data.y_cell.values + + source_obj = np.full(len(x), -1, dtype=int) + for i_obj, obj in enumerate(cluster.objects): + source_obj[obj.mask] = i_obj + + rows, cols = _mst_edges(x, y) + + fig, ax = plt.subplots() + + for r, c in zip(rows, cols): + if source_obj[r] == source_obj[c] and source_obj[r] >= 0: + color = _OBJECT_COLORS[source_obj[r] % len(_OBJECT_COLORS)] + ax.plot([x[r], x[c]], [y[r], y[c]], color=color, linewidth=1.0, zorder=1) + else: + ax.plot( + [x[r], x[c]], [y[r], y[c]], + color="grey", linewidth=0.8, linestyle="--", alpha=0.5, zorder=1, + ) + + for i_obj, obj in enumerate(cluster.objects): + color = _OBJECT_COLORS[i_obj % len(_OBJECT_COLORS)] + ax.scatter(x[obj.mask], y[obj.mask], c=color, label=f"object {i_obj}", zorder=2) + ax.scatter(obj.x_cent, obj.y_cent, marker="+", c=color, s=100, zorder=3) + + ax.set_xlabel("x [cell pixels]") + ax.set_ylabel("y [cell pixels]") + ax.legend(loc="best", fontsize="small") + return fig From 60f752f7fda2d6d668a932a458c5635e947bbeb2 Mon Sep 17 00:00:00 2001 From: Eric Charles Date: Tue, 3 Nov 2026 17:01:52 -0800 Subject: [PATCH 5/5] Add VizMatchesMST.ipynb using MST visualization functions Counterpart to VizMatches.ipynb that loads MST match outputs and displays clusters and objects with showMSTCluster / showMSTObjects instead of the pixel counts-map based showCluster. Generated with AI Co-Authored-By: SLAC AI --- examples/rubin_dp2/VizMatchesMST.ipynb | 218 +++++++++++++++++++++++++ 1 file changed, 218 insertions(+) create mode 100644 examples/rubin_dp2/VizMatchesMST.ipynb diff --git a/examples/rubin_dp2/VizMatchesMST.ipynb b/examples/rubin_dp2/VizMatchesMST.ipynb new file mode 100644 index 0000000..0ed9386 --- /dev/null +++ b/examples/rubin_dp2/VizMatchesMST.ipynb @@ -0,0 +1,218 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "580bf251", + "metadata": {}, + "source": [ + "# Visualize MST matches\n", + "\n", + "Loads saved MST match objects and displays cluster structure using the\n", + "MST-based visualization utilities.\n", + "\n", + "`showMSTCluster` draws the minimum spanning tree of sources in a cluster,\n", + "coloured by catalog ID.\n", + "\n", + "`showMSTObjects` draws the same tree with intra-object edges coloured\n", + "solid by object and inter-object (cut) edges dashed grey, making the\n", + "splitting decisions visible." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd06f0e5-b79f-41bf-ac0a-77f6ba4ac6fb", + "metadata": {}, + "outputs": [], + "source": [ + "# Standard imports\n", + "import tables_io\n", + "import hpmcm\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np" + ] + }, + { + "cell_type": "markdown", + "id": "c2221fae", + "metadata": {}, + "source": [ + "## Configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4b02954e", + "metadata": {}, + "outputs": [], + "source": [ + "tract = 8006" + ] + }, + { + "cell_type": "markdown", + "id": "a1b2c3d4", + "metadata": {}, + "source": [ + "## Rubin shear match" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "63938c5e-0a69-447a-9357-7cfd74aea22b", + "metadata": {}, + "outputs": [], + "source": [ + "# Load saved Rubin shear MST match\n", + "rubin_shear_match = hpmcm.MSTShearMatch.load(\n", + " f\"{tract}/rubin_mst_md_match\",\n", + " (100, 101),\n", + " (100, 101),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a4f422af-9a29-49f1-9b1c-250021d6c763", + "metadata": {}, + "outputs": [], + "source": [ + "# Cluster view: sources coloured by catalog ID, MST edges in grey\n", + "cell_data = rubin_shear_match.cell_dict[20200]\n", + "cluster_id = list(cell_data.cluster_dict.keys())[0]\n", + "cluster_data = cell_data.cluster_dict[cluster_id]\n", + "_ = hpmcm.viz_utils.showMSTCluster(cluster_data, cell_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b1c2d3e4-f5a6-7890-bcde-f01234567891", + "metadata": {}, + "outputs": [], + "source": [ + "# Object view: intra-object edges coloured by object, cut edges dashed grey\n", + "_ = hpmcm.viz_utils.showMSTObjects(cluster_data, cell_data)" + ] + }, + { + "cell_type": "markdown", + "id": "d2e3f4a5", + "metadata": {}, + "source": [ + "## DESC shear match" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "283bc883-3fa3-4066-a98a-a8f76f1cb173", + "metadata": {}, + "outputs": [], + "source": [ + "# Load saved DESC shear MST match\n", + "desc_shear_match = hpmcm.MSTShearMatch.load(\n", + " f\"{tract}/desc_mst_md_match\",\n", + " (100, 101),\n", + " (100, 101),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "88daa9d6-6c97-4796-a05d-fa0fb678ce1d", + "metadata": {}, + "outputs": [], + "source": [ + "# Cluster view\n", + "cell_data = desc_shear_match.cell_dict[20200]\n", + "cluster_id = list(cell_data.cluster_dict.keys())[4]\n", + "cluster_data = cell_data.cluster_dict[cluster_id]\n", + "_ = hpmcm.viz_utils.showMSTCluster(cluster_data, cell_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c3d4e5f6-a7b8-9012-cdef-012345678902", + "metadata": {}, + "outputs": [], + "source": [ + "# Object view\n", + "_ = hpmcm.viz_utils.showMSTObjects(cluster_data, cell_data)" + ] + }, + { + "cell_type": "markdown", + "id": "e4f5a6b7", + "metadata": {}, + "source": [ + "## WCS object match" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "848b7d70-bd4f-45dd-ab48-56e119b99ef3", + "metadata": {}, + "outputs": [], + "source": [ + "# Load saved WCS object MST match\n", + "wcs_match = hpmcm.MSTWcsMatch.load(\n", + " f\"{tract}/all_mst_match\",\n", + " (100, 101),\n", + " (100, 101),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "609f9ba7-aad2-456f-aa39-201b380be50e", + "metadata": {}, + "outputs": [], + "source": [ + "# Cluster view\n", + "cell_data = wcs_match.cell_dict[20200]\n", + "cluster_id = list(cell_data.cluster_dict.keys())[0]\n", + "cluster_data = cell_data.cluster_dict[cluster_id]\n", + "_ = hpmcm.viz_utils.showMSTCluster(cluster_data, cell_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d4e5f6a7-b8c9-0123-defa-123456789013", + "metadata": {}, + "outputs": [], + "source": [ + "# Object view\n", + "_ = hpmcm.viz_utils.showMSTObjects(cluster_data, cell_data)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbformat_minor": 5, + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}