Coverage for python/lsst/pipe/tasks/match_tract_catalog.py: 40%
225 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-23 09:30 +0000
1# This file is part of pipe_tasks.
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__all__ = [
23 'MatchTractCatalogSubConfig', 'MatchTractCatalogSubTask',
24 'MatchTractCatalogConfig', 'MatchTractCatalogTask'
25]
27from abc import ABC, abstractmethod
29import astropy.table
30import numpy as np
31import pandas as pd
34import lsst.afw.geom as afwGeom
35import lsst.geom
36import lsst.pex.config as pexConfig
37import lsst.pipe.base as pipeBase
38import lsst.pipe.base.connectionTypes as cT
39from lsst.geom import SpherePoint
40from lsst.obs.base.utils import TableVStack
41from lsst.skymap import BaseSkyMap
43from .diff_matched_tract_catalog import DiffMatchedTractCatalogTaskBase
45MatchTractCatalogBaseTemplates = {
46 "name_input_cat_ref": "truth_summary",
47 "name_input_cat_target": "objectTable_tract",
48}
51class MatchTractCatalogConnections(
52 pipeBase.PipelineTaskConnections,
53 dimensions=("tract", "skymap"),
54 defaultTemplates=MatchTractCatalogBaseTemplates,
55):
56 cat_ref = cT.Input(
57 doc="Reference object catalog to match from",
58 name="{name_input_cat_ref}",
59 storageClass="ArrowAstropy",
60 dimensions=("tract", "skymap"),
61 deferLoad=True,
62 )
63 cat_target = cT.Input(
64 doc="Target object catalog to match",
65 name="{name_input_cat_target}",
66 storageClass="ArrowAstropy",
67 dimensions=("tract", "skymap"),
68 deferLoad=True,
69 )
70 skymap = cT.Input(
71 doc="Input definition of geometry/bbox and projection/wcs for coadded exposures",
72 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
73 storageClass="SkyMap",
74 dimensions=("skymap",),
75 )
76 cat_output_matched = cT.Output(
77 doc="Target matched catalog with indices of reference matches",
78 name="matched_{name_input_cat_ref}_{name_input_cat_target}",
79 storageClass="ArrowAstropy",
80 dimensions=("tract", "skymap"),
81 )
82 cat_output_ref = cT.Output(
83 doc="Reference matched catalog with indices of target matches",
84 name="match_ref_{name_input_cat_ref}_{name_input_cat_target}",
85 storageClass="ArrowAstropy",
86 dimensions=("tract", "skymap"),
87 )
88 cat_output_target = cT.Output(
89 doc="Target matched catalog with indices of reference matches",
90 name="match_target_{name_input_cat_ref}_{name_input_cat_target}",
91 storageClass="ArrowAstropy",
92 dimensions=("tract", "skymap"),
93 )
95 def __init__(self, *, config=None):
96 if not config.output_matched_catalog:
97 del self.cat_output_matched
98 if config.match_multiple_target:
99 # allConnections will change during iteration
100 for name, connection in tuple(self.allConnections.items()):
101 if connection.storageClass == "SkyMap":
102 continue
103 kwargs = {"deferLoad": connection.deferLoad} if hasattr(connection, "deferLoad") else {}
104 delattr(self, name)
105 setattr(
106 self,
107 name,
108 type(connection)(
109 doc=connection.doc,
110 name=connection.name,
111 storageClass=connection.storageClass,
112 dimensions=connection.dimensions,
113 multiple=True,
114 **kwargs
115 ),
116 )
117 self.dimensions = {"skymap"}
118 is_refcat_sharding_none = False
119 if config.refcat_sharding_type != "tract":
120 if config.refcat_sharding_type == "none":
121 is_refcat_sharding_none = True
122 old = self.cat_ref
123 self.cat_ref = cT.Input(
124 doc=old.doc,
125 name=old.name,
126 storageClass=old.storageClass,
127 dimensions=(),
128 deferLoad=old.deferLoad,
129 )
130 else:
131 raise NotImplementedError(f"{config.refcat_sharding_type=} not implemented")
132 if config.target_sharding_type != "tract":
133 if config.target_sharding_type == "none":
134 if config.match_multiple_target:
135 raise ValueError(
136 f"Incompatible {config.match_multiple_target=} with {config.target_sharding_type=}"
137 )
138 old = self.cat_target
139 self.cat_target = cT.Input(
140 doc=old.doc,
141 name=old.name,
142 storageClass=old.storageClass,
143 dimensions=(),
144 deferLoad=old.deferLoad,
145 )
146 if is_refcat_sharding_none:
147 for suffix_old in ("matched", "ref", "target"):
148 name_old = f"cat_output_{suffix_old}"
149 old = getattr(self, name_old)
150 setattr(
151 self,
152 name_old,
153 cT.Output(
154 doc=old.doc,
155 name=old.name,
156 storageClass=old.storageClass,
157 dimensions=(),
158 ),
159 )
160 del self.skymap
161 self.dimensions = set()
162 else:
163 raise NotImplementedError(f"{config.target_sharding_type=} not implemented")
166class MatchTractCatalogSubConfig(pexConfig.Config):
167 """Config class for the MatchTractCatalogSubTask to define methods returning
168 values that depend on multiple config settings.
169 """
170 @property
171 @abstractmethod
172 def columns_in_ref(self) -> set[str]:
173 raise NotImplementedError()
175 @property
176 def columns_ordered_in_ref(self) -> dict[str, None]:
177 raise NotImplementedError()
179 @property
180 @abstractmethod
181 def columns_in_target(self) -> set[str]:
182 raise NotImplementedError()
184 @property
185 def columns_ordered_in_target(self) -> dict[str, None]:
186 raise NotImplementedError()
189class MatchTractCatalogSubTask(pipeBase.Task, ABC):
190 """An abstract interface for subtasks of MatchTractCatalogTask to match
191 two tract object catalogs.
193 Parameters
194 ----------
195 **kwargs
196 Additional arguments to be passed to the `lsst.pipe.base.Task`
197 constructor.
198 """
199 ConfigClass = MatchTractCatalogSubConfig
201 def __init__(self, **kwargs):
202 super().__init__(**kwargs)
204 @abstractmethod
205 def run(
206 self,
207 catalog_ref: pd.DataFrame | astropy.table.Table,
208 catalog_target: pd.DataFrame | astropy.table.Table,
209 wcs: afwGeom.SkyWcs = None,
210 ) -> pipeBase.Struct:
211 """Match sources in a reference tract catalog with a target catalog.
213 Parameters
214 ----------
215 catalog_ref : `pandas.DataFrame` | `astropy.table.Table`
216 A reference catalog to match objects/sources from.
217 catalog_target : `pandas.DataFrame` | `astropy.table.Table`
218 A target catalog to match reference objects/sources to.
219 wcs : `lsst.afw.image.SkyWcs`
220 A coordinate system to convert catalog positions to sky coordinates.
222 Returns
223 -------
224 retStruct : `lsst.pipe.base.Struct`
225 A struct with output_ref and output_target attribute containing the
226 output matched catalogs.
227 """
228 raise NotImplementedError()
231class MatchTractCatalogConfig(
232 pipeBase.PipelineTaskConfig,
233 pipelineConnections=MatchTractCatalogConnections,
234):
235 """Configure a MatchTractCatalogTask, including a configurable matching subtask.
236 """
237 coord_unit = pexConfig.Field[str](
238 doc="lsst.geom unit (or astropy equivalent) of the coordinate columns."
239 "Only used to determine the tract for rows in non-tract-sharded "
240 "catalogs without a tract column.",
241 optional=True,
242 )
243 diff_matched_catalog = pexConfig.ConfigurableField(
244 target=DiffMatchedTractCatalogTaskBase,
245 doc="Task to make a matched catalog out of the match index tables",
246 )
247 match_multiple_target = pexConfig.Field[bool](
248 doc="Whether to match multiple target tract catalogs",
249 default=False,
250 )
251 match_tract_catalog = pexConfig.ConfigurableField(
252 target=MatchTractCatalogSubTask,
253 doc="Task to match sources in a reference tract catalog with a target catalog",
254 )
255 output_matched_catalog = pexConfig.Field[bool](
256 doc="Whether to run the diff_matched_catalog task and write a matched catalog,"
257 " not just the catalogs of match indices",
258 default=False,
259 )
260 refcat_sharding_type = pexConfig.ChoiceField[str](
261 doc="The type of sharding (spatial splitting) for the reference catalog",
262 allowed={"tract": "Tract-based shards", "none": "No sharding at all"},
263 default="tract",
264 )
265 target_sharding_type = pexConfig.ChoiceField[str](
266 doc="The type of sharding (spatial splitting) for the target catalog",
267 allowed={"tract": "Tract-based shards", "none": "No sharding at all"},
268 default="tract",
269 )
271 def get_columns_in(self) -> tuple[set, set]:
272 """Get the set of input columns required for matching.
274 This function exists for backward compatibility and simply returns
275 the results of get_columns_ordered_in cast to sets.
276 """
277 columns_ref, columns_target = self.get_columns_ordered_in()
278 return set(columns_ref.keys()), set(columns_target.keys())
280 def get_columns_ordered_in(self) -> tuple[dict[str, None], dict[str, None]]:
281 """Get the ordered set of input columns required for matching.
283 Returns
284 -------
285 columns_ref : `set` [`str`]
286 The set of required input catalog column names.
287 columns_target : `set` [`str`]
288 The set of required target catalog column names.
289 """
290 try:
291 columns_ref, columns_target = (
292 self.match_tract_catalog.columns_ordered_in_ref,
293 self.match_tract_catalog.columns_ordered_in_target,
294 )
295 except AttributeError as err:
296 raise RuntimeError(f'{__class__}.match_tract_catalog must have columns_in_ref and'
297 f' columns_in_target attributes: {err}') from None
298 if self.output_matched_catalog:
299 config_diff = self.diff_matched_catalog.value
300 columns_ref.update({k: None for k in config_diff.columns_in_ref})
301 columns_target.update({k: None for k in config_diff.columns_in_target})
302 return columns_ref, columns_target
305class MatchTractCatalogTask(pipeBase.PipelineTask):
306 """Match sources in a reference tract catalog with those in a target catalog.
307 """
308 ConfigClass = MatchTractCatalogConfig
309 _DefaultName = "MatchTractCatalog"
311 def __init__(self, initInputs, **kwargs):
312 super().__init__(initInputs=initInputs, **kwargs)
313 self.makeSubtask("match_tract_catalog")
314 self.makeSubtask("diff_matched_catalog")
316 _astropy_u_to_lsst_geom = {
317 "arcmin": "arcminutes",
318 "deg": "degrees",
319 "rad": "radians",
320 }
322 def runQuantum(self, butlerQC, inputRefs, outputRefs):
323 inputs = butlerQC.get(inputRefs)
324 columns_ref, columns_target = (tuple(columns) for columns in self.config.get_columns_ordered_in())
325 skymap = inputs.pop("skymap") if (self.config.refcat_sharding_type != 'none') or (
326 self.config.target_sharding_type != 'none') else None
327 is_refcat_per_tract = self.config.refcat_sharding_type == 'tract'
328 is_target_per_tract = self.config.target_sharding_type == 'tract'
330 if self.config.match_multiple_target:
331 if is_target_per_tract:
332 names_columns = [("cat_target", columns_target)]
333 catalogs = {}
334 else:
335 catalogs = {"cat_target": inputs["cat_target"].get(parameters={'columns': columns_target})}
336 names_columns = []
337 if is_refcat_per_tract:
338 names_columns.append(("cat_ref", columns_ref))
339 else:
340 catalogs["cat_ref"] = inputs["cat_ref"].get(parameters={'columns': columns_ref})
342 for name, columns in names_columns:
343 extra_values = {}
344 handles = []
345 for idx, (tract, handle) in enumerate(sorted(
346 (inputRef.dataId["tract"], inputHandle)
347 for inputRef, inputHandle in zip(getattr(inputRefs, name), inputs[name], strict=True)
348 )):
349 handles.append(handle)
350 extra_values[idx] = {"tract": tract}
351 catalogs[name] = TableVStack.vstack_handles(
352 handles,
353 extra_values=extra_values,
354 kwargs_get={"parameters": {"columns": columns}}
355 )
356 catalog_ref, catalog_target = catalogs["cat_ref"], catalogs["cat_target"]
357 else:
358 catalog_ref, catalog_target = (
359 inputs[name].get(parameters={'columns': columns})
360 for name, columns in (('cat_ref', columns_ref), ('cat_target', columns_target))
361 )
362 if self.config.match_multiple_target:
363 self._add_tract_column_to_catalogs(catalog_ref, catalog_target, skymap)
365 outputs = self.run(
366 catalog_ref=catalog_ref,
367 catalog_target=catalog_target,
368 wcs=None if (self.config.match_multiple_target or (skymap is None)) else (
369 skymap[butlerQC.quantum.dataId["tract"]].wcs),
370 )
372 if self.config.match_multiple_target:
373 outputs_new = {}
375 for name, catalog_in, name_tract in (
376 ("cat_output_ref", catalog_ref, outputs.name_tract_ref_in),
377 ("cat_output_target", catalog_target, "tract"),
378 ):
379 catalogs_out = []
380 for outputRef in getattr(outputRefs, name):
381 tract = outputRef.dataId["tract"]
382 catalog_out = getattr(outputs, name)
383 catalogs_out.append(catalog_out[catalog_in[name_tract] == tract])
384 outputs_new[name] = catalogs_out
386 if self.config.output_matched_catalog:
387 cat_matched = outputs.cat_output_matched
388 catalogs_out = []
389 tract_target = cat_matched[outputs.name_tract_target]
390 patch_target = cat_matched[outputs.name_patch_target]
391 masked_target = tract_target.mask == True # noqa: E712
393 tract_ref = cat_matched[outputs.name_tract_ref]
394 tract_out = np.array(tract_target)
395 tract_out[masked_target] = tract_ref[masked_target]
396 cat_matched["tract"] = tract_out
397 cat_matched["tract"].description = (f"{outputs.name_tract_target} if available"
398 f" else {outputs.name_tract_ref}")
399 patch_ref = cat_matched[outputs.name_patch_ref]
400 patch_out = np.array(patch_target)
401 # The patch and target masks should be identical
402 patch_out[masked_target] = patch_ref[masked_target]
403 cat_matched["patch"] = patch_out
404 cat_matched["patch"].description = (f"{outputs.name_patch_target} if available"
405 f" else {outputs.name_patch_ref}")
407 for outputRef in outputRefs.cat_output_matched:
408 tract = outputRef.dataId["tract"]
409 catalogs_out.append(cat_matched[tract_out == tract])
410 outputs_new["cat_output_matched"] = catalogs_out
412 outputs = pipeBase.Struct(
413 **outputs_new,
414 **{k: v for k, v in outputs.getDict().items() if k not in outputs_new},
415 )
417 butlerQC.put(outputs, outputRefs)
419 def run(
420 self,
421 catalog_ref: pd.DataFrame | astropy.table.Table,
422 catalog_target: pd.DataFrame | astropy.table.Table,
423 wcs: afwGeom.SkyWcs = None,
424 ) -> pipeBase.Struct:
425 """Match sources in a reference tract catalog with a target catalog.
427 Parameters
428 ----------
429 catalog_ref : `pandas.DataFrame` | `astropy.table.Table`
430 A reference catalog to match objects/sources from.
431 catalog_target : `pandas.DataFrame` | `astropy.table.Table`
432 A target catalog to match reference objects/sources to.
433 wcs : `lsst.afw.image.SkyWcs`
434 A coordinate system to convert catalog positions to sky coordinates,
435 if necessary.
437 Returns
438 -------
439 retStruct : `lsst.pipe.base.Struct`
440 A struct with output_ref and output_target attribute containing the
441 output matched catalogs. If self.config.output_matched_catalog is
442 True, cat_output_matched is also returned, along with the names of
443 the tract/patch columns, which may have been changed to avoid
444 collisions.
445 """
446 output = self.match_tract_catalog.run(catalog_ref, catalog_target, wcs=wcs)
447 if output.exceptions: 447 ↛ 448line 447 didn't jump to line 448 because the condition on line 447 was never true
448 self.log.warning('Exceptions: %s', output.exceptions)
449 retStruct = pipeBase.Struct(cat_output_ref=output.cat_output_ref,
450 cat_output_target=output.cat_output_target)
452 if self.config.output_matched_catalog: 452 ↛ 491line 452 didn't jump to line 491 because the condition on line 452 was always true
453 # TODO: Consider making this a config parameter
454 # Making it optional is probably preferable than quietly
455 # checking a hardcoded column name
456 name_tract_ref_in = "tract"
457 diff_prefix_ref = self.config.diff_matched_catalog.value.column_matched_prefix_ref
458 diff_prefix_target = self.config.diff_matched_catalog.value.column_matched_prefix_target
459 name_tract_ref = f"{diff_prefix_ref}tract"
460 name_tract_target = f"{diff_prefix_target}tract"
461 name_patch_ref = f"{diff_prefix_ref}patch"
462 name_patch_target = f"{diff_prefix_target}patch"
463 # Avoid collisions if, for example, both prefixes are ''
464 if (name_tract_ref_in in catalog_ref.colnames) and (name_tract_ref == name_tract_target): 464 ↛ 465line 464 didn't jump to line 465 because the condition on line 464 was never true
465 prefix_new = "ref_" if diff_prefix_ref != "ref_" else "refcat_"
466 name_new = f"{prefix_new}tract"
467 catalog_ref.rename_column(name_tract_ref, name_new)
468 name_tract_ref = name_new
469 name_tract_ref_in = name_new
470 if name_patch_ref in catalog_ref.colnames:
471 name_new = f"{prefix_new}patch"
472 catalog_ref.rename_column(name_patch_ref, name_new)
473 name_patch_ref = name_new
475 outputs_new = self.diff_matched_catalog.run(
476 catalog_ref=catalog_ref,
477 catalog_target=catalog_target,
478 catalog_match_ref=retStruct.cat_output_ref,
479 catalog_match_target=retStruct.cat_output_target,
480 )
481 retStruct = pipeBase.Struct(
482 **retStruct.getDict(),
483 cat_output_matched=outputs_new.cat_matched,
484 name_tract_ref=name_tract_ref,
485 name_tract_ref_in=name_tract_ref_in,
486 name_tract_target=name_tract_target,
487 name_patch_ref=name_patch_ref,
488 name_patch_target=name_patch_target,
489 )
491 return retStruct
493 def _add_tract_column_to_catalogs(self, catalog_ref, catalog_target, skymap, add_patch: bool = True):
494 """Add a tract column to catalogs that may be missing it.
496 Parameters
497 ----------
498 catalog_ref
499 A reference catalog.
500 catalog_target
501 A target catalog.
502 skymap
503 The skymap info to use.
504 add_patch
505 Whether to add a patch column as well.
506 """
507 errors = []
508 if compute_tract_target := ("tract" not in catalog_target.colnames): 508 ↛ 513line 508 didn't jump to line 513 because the condition on line 508 was always true
509 if self.config.target_sharding_type != "none": 509 ↛ 510line 509 didn't jump to line 510 because the condition on line 509 was never true
510 errors.append(
511 f"Target catalog has no tract column with {self.config.target_sharding_type=} != 'none'"
512 )
513 if compute_tract_ref := ("tract" not in catalog_ref.colnames): 513 ↛ 518line 513 didn't jump to line 518 because the condition on line 513 was always true
514 if self.config.refcat_sharding_type != "none": 514 ↛ 515line 514 didn't jump to line 515 because the condition on line 514 was never true
515 errors.append(
516 f"Ref catalog has no tract column with {self.config.refcat_sharding_type=} != 'none'"
517 )
518 if errors: 518 ↛ 519line 518 didn't jump to line 519 because the condition on line 518 was never true
519 raise RuntimeError("; ".join(errors))
520 compute_patch_ref = add_patch and ("patch" not in catalog_ref.colnames)
521 compute_patch_target = add_patch and ("patch" not in catalog_target.colnames)
522 unit_dict = self._astropy_u_to_lsst_geom.copy()
523 if compute_tract_target or compute_tract_ref or compute_patch_target or compute_patch_ref: 523 ↛ exitline 523 didn't return from function '_add_tract_column_to_catalogs' because the condition on line 523 was always true
524 if not self.config.diff_matched_catalog.coord_format.coords_spherical: 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true
525 raise RuntimeError(
526 f"Can't compute tract columns with unless"
527 f" {self.config.diff_matched_catalog.coord_format.coords_spherical} == True"
528 )
529 ref_c, target_c = self.config.diff_matched_catalog.coord_format.format_catalogs(
530 catalog_ref=catalog_ref, catalog_target=catalog_target,
531 )
532 cats_add = []
533 if compute_tract_target or compute_patch_target: 533 ↛ 537line 533 didn't jump to line 537 because the condition on line 533 was always true
534 cats_add.append((
535 target_c, catalog_target, "target", compute_tract_target, compute_patch_target
536 ))
537 if compute_tract_ref: 537 ↛ 539line 537 didn't jump to line 539 because the condition on line 537 was always true
538 cats_add.append((ref_c, catalog_ref, "ref", compute_tract_ref, compute_patch_ref))
539 for cat_c, catalog_add, name_c, compute_tract, compute_patch in cats_add:
540 if (unit := getattr(catalog_add[cat_c.column_coord1], "unit")) is None: 540 ↛ 549line 540 didn't jump to line 549 because the condition on line 540 was always true
541 unit = self.config.coord_unit
542 if unit is None: 542 ↛ 543line 542 didn't jump to line 543 because the condition on line 542 was never true
543 raise RuntimeError(
544 f"Must specify coord_unit since {name_c} column={cat_c.column_coord1}"
545 f" has no units"
546 )
547 unit = getattr(lsst.geom, unit, getattr(lsst.geom, unit_dict[str(unit)]))
548 else:
549 unit = getattr(lsst.geom, unit_dict[str(unit)])
550 coords = [SpherePoint(ra, dec, unit) for ra, dec in zip(cat_c.coord1, cat_c.coord2)]
551 if compute_tract: 551 ↛ 553line 551 didn't jump to line 553 because the condition on line 551 was always true
552 catalog_add["tract"] = np.array([skymap.findTract(coord).getId() for coord in coords])
553 if compute_patch: 553 ↛ 539line 553 didn't jump to line 539 because the condition on line 553 was always true
554 tracts = catalog_add["tract"]
555 patches = np.array([
556 skymap[tract].findPatch(coord).getSequentialIndex()
557 for coord, tract in zip(coords, tracts)
558 ])
559 catalog_add["patch"] = patches