Coverage for python/lsst/analysis/ap/compare.py: 78%
151 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-24 10:06 +0000
1# This file is part of analysis_ap.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22"""Catalog cross-matching and pair-wise comparison utilities for
23DiaSource-like tables.
24"""
26__all__ = ["match_catalogs", "flux_residuals", "match_to_truth"]
28import numpy as np
29import pandas as pd
31import astropy.units as u
32from astropy.coordinates import SkyCoord, match_coordinates_sky, search_around_sky
35def _as_arcsec(radius):
36 """Coerce a number-or-Quantity radius to a float in arcseconds."""
37 if isinstance(radius, u.Quantity):
38 return float(radius.to_value(u.arcsec))
39 return float(radius)
42def _group_keys(srcs1, srcs2, on):
43 """Yield the union of distinct ``on``-tuples present in either frame."""
44 if not on:
45 yield None
46 return
47 s1 = set(map(tuple, srcs1[list(on)].itertuples(index=False, name=None)))
48 s2 = set(map(tuple, srcs2[list(on)].itertuples(index=False, name=None)))
49 for key in sorted(s1 | s2):
50 yield key
53def _select_group(df, on, key):
54 """Return the slice of ``df`` whose ``on`` columns equal ``key``."""
55 if key is None or not on:
56 return df
57 mask = pd.Series(True, index=df.index)
58 for col, val in zip(on, key):
59 mask &= df[col] == val
60 return df[mask]
63def _one_to_one_pairs(coords1, coords2, rad_arcsec):
64 """Pair two sets of positions one to one, closest pairs first.
66 Parameters
67 ----------
68 coords1, coords2 : `astropy.coordinates.SkyCoord`
69 Positions to pair; both non-empty arrays.
70 rad_arcsec : `float`
71 Maximum separation of a pair, in arcseconds.
73 Returns
74 -------
75 index1, index2 : `numpy.ndarray`
76 Indices into ``coords1`` and ``coords2`` of each pair.
77 sep_arcsec : `numpy.ndarray`
78 Separation of each pair, in arcseconds.
79 """
80 idx1, idx2, sep, _ = search_around_sky(coords1, coords2, rad_arcsec*u.arcsec)
81 sep = sep.to_value(u.arcsec)
82 # Ties are broken by index so the result does not depend on the order
83 # search_around_sky returns pairs in.
84 order = np.lexsort((idx2, idx1, sep))
85 used1, used2 = set(), set()
86 keep = []
87 for k in order:
88 if idx1[k] not in used1 and idx2[k] not in used2:
89 used1.add(idx1[k])
90 used2.add(idx2[k])
91 keep.append(k)
92 keep = np.array(keep, dtype=int)
93 return idx1[keep], idx2[keep], sep[keep]
96def match_catalogs(srcs1, srcs2, radius=0.5*u.arcsec, on=("visit", "detector"),
97 ra_col="ra", dec_col="dec", id_col="diaSourceId", one_to_one=False):
98 """Spatially cross-match two DiaSource-like DataFrames.
100 Sources are first partitioned by the columns in
101 ``on`` (e.g. matched only within the same (visit, detector)), then matched
102 via nearest-neighbor on the sphere.
104 Parameters
105 ----------
106 srcs1, srcs2 : `pandas.DataFrame`
107 Source tables. Each must contain ``ra_col``, ``dec_col``, ``id_col``,
108 and every column listed in ``on``.
109 radius : `astropy.units.Quantity` or `float`
110 Maximum separation for a pair to count as matched. A bare float is
111 interpreted as arcseconds.
112 on : `tuple` [`str`]
113 Columns to group on before matching. Pass an empty tuple to match the
114 full catalog with no grouping.
115 ra_col, dec_col : `str`
116 Column names for sky coordinates, in degrees.
117 id_col : `str`
118 Column name for the source id in both catalogs.
119 one_to_one : `bool`, optional
120 If True, each source is in at most one pair: pairs within ``radius``
121 are taken in order of increasing separation, skipping any whose
122 sources are already paired. If False, each source in ``srcs1`` is
123 paired with its nearest neighbor in ``srcs2``, so several may share
124 one partner.
126 Returns
127 -------
128 matched : `pandas.DataFrame`
129 Rows from ``srcs1`` that found a partner in ``srcs2`` within
130 ``radius``. Two columns are added: ``<id_col>_2`` with the partner's
131 id, and ``xmatch_dist_arcsec`` with the on-sky separation in arcsec.
132 unique1 : `pandas.DataFrame`
133 Rows from ``srcs1`` with no partner.
134 unique2 : `pandas.DataFrame`
135 Rows from ``srcs2`` not pointed at by any matched pair.
136 """
137 rad_arcsec = _as_arcsec(radius)
138 on = tuple(on)
139 id2_col = f"{id_col}_2"
141 matched_chunks = []
142 unique1_chunks = []
143 unique2_chunks = []
145 for key in _group_keys(srcs1, srcs2, on):
146 gs1 = _select_group(srcs1, on, key).copy()
147 gs2 = _select_group(srcs2, on, key).copy()
149 if len(gs1) == 0:
150 unique2_chunks.append(gs2)
151 continue
152 if len(gs2) == 0:
153 unique1_chunks.append(gs1)
154 continue
156 coords1 = SkyCoord(ra=gs1[ra_col].values*u.deg,
157 dec=gs1[dec_col].values*u.deg)
158 coords2 = SkyCoord(ra=gs2[ra_col].values*u.deg,
159 dec=gs2[dec_col].values*u.deg)
160 if one_to_one:
161 idx1, idx2, sep = _one_to_one_pairs(coords1, coords2, rad_arcsec)
162 gs1["xmatch_dist_arcsec"] = np.nan
163 gs1[id2_col] = pd.Series(pd.NA, index=gs1.index, dtype=object)
164 gs1.iloc[idx1, gs1.columns.get_loc("xmatch_dist_arcsec")] = sep
165 gs1.iloc[idx1, gs1.columns.get_loc(id2_col)] = gs2[id_col].values[idx2]
166 else:
167 idx, sep, _ = match_coordinates_sky(coords1, coords2)
168 gs1["xmatch_dist_arcsec"] = sep.to_value(u.arcsec)
169 gs1[id2_col] = gs2[id_col].values[idx]
171 has_match = gs1["xmatch_dist_arcsec"] <= rad_arcsec
172 m = gs1[has_match].copy()
173 m[id2_col] = m[id2_col].astype(gs2[id_col].dtype)
174 u1 = gs1[~has_match].drop(columns=["xmatch_dist_arcsec", id2_col])
175 u2 = gs2[~gs2[id_col].isin(set(m[id2_col]))]
177 matched_chunks.append(m)
178 unique1_chunks.append(u1)
179 unique2_chunks.append(u2)
181 matched = (pd.concat(matched_chunks) if matched_chunks
182 else srcs1.iloc[0:0].assign(**{"xmatch_dist_arcsec": np.nan,
183 id2_col: pd.NA}))
184 unique1 = pd.concat(unique1_chunks) if unique1_chunks else srcs1.iloc[0:0].copy()
185 unique2 = pd.concat(unique2_chunks) if unique2_chunks else srcs2.iloc[0:0].copy()
186 return matched, unique1, unique2
189def flux_residuals(matched, srcs2, flux_col="psfFlux", err_col="psfFluxErr",
190 id_col="diaSourceId", plot=False):
191 """Compute per-pair flux residuals from a `match_catalogs` result.
193 Parameters
194 ----------
195 matched : `pandas.DataFrame`
196 Output of `match_catalogs` (rows from catalog 1 with partner ids).
197 srcs2 : `pandas.DataFrame`
198 Catalog 2, indexed implicitly by ``id_col`` for partner lookup.
199 flux_col, err_col : `str`
200 Column names for the flux and its error in both catalogs.
201 id_col : `str`
202 Source-id column in both catalogs. ``matched`` is assumed to have
203 ``f"{id_col}_2"`` populated by `match_catalogs`.
204 plot : `bool`
205 If True, return a histogram + Q-Q plot in addition to the residuals.
207 Returns
208 -------
209 residuals : `pandas.DataFrame`
210 One row per pair with columns ``flux1``, ``flux2``, ``err1``, ``err2``,
211 ``delta_flux``, ``delta_flux_sigma``, and ``xmatch_dist_arcsec``.
212 fig : `matplotlib.figure.Figure`, optional
213 Only returned when ``plot=True``.
214 """
215 id2_col = f"{id_col}_2"
216 s2 = srcs2.set_index(id_col)
217 partner_ids = matched[id2_col].values
219 f1 = matched[flux_col].to_numpy(dtype=float, copy=True)
220 e1 = matched[err_col].to_numpy(dtype=float, copy=True)
221 f2 = s2.loc[partner_ids, flux_col].to_numpy(dtype=float, copy=True)
222 e2 = s2.loc[partner_ids, err_col].to_numpy(dtype=float, copy=True)
224 delta = f1 - f2
225 sigma = np.sqrt(e1**2 + e2**2)
226 with np.errstate(divide="ignore", invalid="ignore"):
227 delta_sigma = np.where(sigma > 0, delta / sigma, np.nan)
229 residuals = pd.DataFrame({
230 id_col: matched[id_col].values,
231 id2_col: partner_ids,
232 "flux1": f1,
233 "flux2": f2,
234 "err1": e1,
235 "err2": e2,
236 "delta_flux": delta,
237 "delta_flux_sigma": delta_sigma,
238 "xmatch_dist_arcsec": matched["xmatch_dist_arcsec"].values,
239 })
241 if not plot: 241 ↛ 243line 241 didn't jump to line 243 because the condition on line 241 was always true
242 return residuals
243 return residuals, _plot_flux_residuals(residuals, flux_col)
246def _plot_flux_residuals(residuals, flux_col):
247 """Histogram + Q-Q plot of normalized flux residuals."""
248 import matplotlib.pyplot as plt
249 sigma = residuals["delta_flux_sigma"].dropna().values
250 if len(sigma) == 0:
251 fig, ax = plt.subplots()
252 ax.text(0.5, 0.5, "no finite residuals", ha="center", va="center")
253 return fig
255 fig, axes = plt.subplots(1, 2, figsize=(10, 4))
256 axes[0].hist(sigma, bins=50, color="C0")
257 axes[0].axvline(0, color="grey", lw=0.5)
258 axes[0].set_xlabel(rf"$\Delta {flux_col} / \sigma$")
259 axes[0].set_ylabel("count")
260 axes[0].set_title(f"N={len(sigma)}, "
261 f"med={np.median(sigma):.3f}, "
262 f"MAD={np.median(np.abs(sigma - np.median(sigma))):.3f}")
264 # Quick Q-Q vs standard normal without depending on scipy.
265 sp = np.sort(sigma)
266 # Inverse-CDF approximation for the standard normal via erfinv.
267 quantiles = (np.arange(len(sp)) + 0.5) / len(sp)
268 expected = np.sqrt(2) * _erfinv(2*quantiles - 1)
269 axes[1].plot(expected, sp, ".", ms=2)
270 lim = max(abs(expected[0]), abs(expected[-1]), 3.0)
271 axes[1].plot([-lim, lim], [-lim, lim], "k--", lw=0.5)
272 axes[1].set_xlabel("expected (N(0,1))")
273 axes[1].set_ylabel("observed")
274 axes[1].set_title("Q-Q vs standard normal")
275 fig.tight_layout()
276 return fig
279def _erfinv(y):
280 """Approximate inverse error function (vectorized).
282 Uses the formula from Winitzki (2008); accurate to ~4e-3 across the
283 domain, which is plenty for plotting Q-Q lines.
284 """
285 a = 0.147
286 sign = np.sign(y)
287 ln1 = np.log(np.clip(1 - y*y, 1e-300, 1.0))
288 term = 2/(np.pi*a) + ln1/2
289 return sign * np.sqrt(np.sqrt(term*term - ln1/a) - term)
292def match_to_truth(srcs, truth, radius=0.5*u.arcsec,
293 src_ra="ra", src_dec="dec", src_id="diaSourceId",
294 truth_ra="ra", truth_dec="dec", truth_id="injection_id"):
295 """Match a detected catalog against a truth/injection catalog.
297 Performs the match in both directions to compute purity (fraction of
298 detections that correspond to a real injected source) and completeness
299 (fraction of injected sources recovered).
301 Parameters
302 ----------
303 srcs : `pandas.DataFrame`
304 Detected sources.
305 truth : `pandas.DataFrame`
306 Truth catalog, e.g. an injection catalog.
307 radius : `astropy.units.Quantity` or `float`
308 Maximum separation for a match.
309 src_ra, src_dec, src_id : `str`
310 Column names in ``srcs``.
311 truth_ra, truth_dec, truth_id : `str`
312 Column names in ``truth``.
314 Returns
315 -------
316 out : `dict`
317 With keys:
319 - ``"srcs"``: a copy of ``srcs`` with three columns appended:
320 ``"is_real"`` (bool), ``"truth_dist_arcsec"`` (float), and
321 ``f"{truth_id}_match"`` (matched truth id, NA if no match).
322 - ``"truth"``: a copy of ``truth`` with three columns appended:
323 ``"detected"``, ``"detection_dist_arcsec"``, and
324 ``f"{src_id}_match"``.
325 - ``"purity"``: float, fraction of ``srcs`` rows with ``is_real``.
326 - ``"completeness"``: float, fraction of ``truth`` rows with
327 ``detected``.
328 """
329 rad_arcsec = _as_arcsec(radius)
330 srcs_out = srcs.copy()
331 truth_out = truth.copy()
333 truth_id_match = f"{truth_id}_match"
334 src_id_match = f"{src_id}_match"
336 if len(srcs_out) == 0 or len(truth_out) == 0: 336 ↛ 337line 336 didn't jump to line 337 because the condition on line 336 was never true
337 srcs_out["is_real"] = False
338 srcs_out["truth_dist_arcsec"] = np.nan
339 srcs_out[truth_id_match] = pd.NA
340 truth_out["detected"] = False
341 truth_out["detection_dist_arcsec"] = np.nan
342 truth_out[src_id_match] = pd.NA
343 return {"srcs": srcs_out, "truth": truth_out,
344 "purity": 0.0, "completeness": 0.0}
346 sc_src = SkyCoord(srcs_out[src_ra].values*u.deg,
347 srcs_out[src_dec].values*u.deg)
348 sc_tru = SkyCoord(truth_out[truth_ra].values*u.deg,
349 truth_out[truth_dec].values*u.deg)
351 idx_to_truth, sep_to_truth, _ = match_coordinates_sky(sc_src, sc_tru)
352 srcs_out["truth_dist_arcsec"] = sep_to_truth.to_value(u.arcsec)
353 srcs_out[truth_id_match] = truth_out[truth_id].values[idx_to_truth]
354 srcs_out["is_real"] = srcs_out["truth_dist_arcsec"] <= rad_arcsec
355 srcs_out.loc[~srcs_out["is_real"], truth_id_match] = pd.NA
357 idx_to_src, sep_to_src, _ = match_coordinates_sky(sc_tru, sc_src)
358 truth_out["detection_dist_arcsec"] = sep_to_src.to_value(u.arcsec)
359 truth_out[src_id_match] = srcs_out[src_id].values[idx_to_src]
360 truth_out["detected"] = truth_out["detection_dist_arcsec"] <= rad_arcsec
361 truth_out.loc[~truth_out["detected"], src_id_match] = pd.NA
363 purity = float(srcs_out["is_real"].sum() / len(srcs_out))
364 completeness = float(truth_out["detected"].sum() / len(truth_out))
365 return {"srcs": srcs_out, "truth": truth_out,
366 "purity": purity, "completeness": completeness}