Coverage for tests/pipeline.py: 50%
138 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-06 08:34 +0000
1# This file is part of meas_extensions_scarlet.
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"""Staged pipeline helpers for the deblender test suite.
24The four stages — :func:`build_image`, :func:`detect`,
25:func:`deconvolve`, :func:`deblend` — are pure functions over
26:class:`Scene` and task-config inputs. Each stage memoizes its
27result keyed by the scene name and the configs that have flowed
28through, so asking for the same (scene, config) tuple in a later
29test method returns the same bundle without re-running the work.
31Tests that want to *mutate* a bundle's contents (e.g. by calling
32``updateCatalogFootprints`` against a catalog) must copy the
33piece they intend to mutate first; the bundles themselves are
34frozen but the LSST objects they reference are not.
35"""
37import json
38from dataclasses import dataclass
39from typing import Any
41import numpy as np
43import lsst.afw.image as afwImage
44import lsst.scarlet.lite as scl
45from lsst.afw.detection import GaussianPsf
46from lsst.afw.table import Schema, SchemaMapper, SourceCatalog, SourceTable
47from lsst.meas.algorithms import SourceDetectionTask
48from lsst.meas.extensions.scarlet.deconvolveExposureTask import DeconvolveExposureTask
49from lsst.meas.extensions.scarlet.scarletDeblendTask import ScarletDeblendTask
50from lsst.pipe.base import Struct
52from scenes import Scene
53from utils import initData
56# Per-band Gaussian PSF sigmas, in pixels, indexed in scene-band order.
57_PSF_SIGMAS: tuple[float, ...] = (1.0, 1.2, 1.4)
58# Half-size of the square PSF kernel in pixels (kernel is 2*r+1 on a side).
59_PSF_RADIUS: int = 20
60# Sigma of the narrow model PSF that scenes are rendered into, in pixels.
61_MODEL_PSF_SIGMA: float = 0.8
62# Peak-to-peak amplitude of the uniform noise added to convolved images.
63_NOISE_AMPLITUDE: float = 0.05
64# RNG seed for the noise draw; fixed so the cache is reproducible.
65_NOISE_SEED: int = 0
68@dataclass(frozen=True)
69class ImageBundle:
70 """Rendered image stage of the pipeline for one scene.
72 Holds both the raw model and the observed (PSF-convolved, noisy)
73 image of a single scene, plus the PSF objects needed by downstream
74 stages.
76 Parameters
77 ----------
78 scene : `Scene`
79 Scene that produced this bundle.
80 deconvolved : `lsst.scarlet.lite.Image`
81 Model image convolved with the narrow model PSF only (the
82 "truth" image scarlet is asked to recover).
83 convolved : `lsst.scarlet.lite.Image`
84 Model image convolved with the full per-band image PSFs;
85 noise-free.
86 noise : `numpy.ndarray`
87 Per-pixel noise realization added to ``convolved`` to produce
88 ``noisy_image``.
89 noisy_image : `lsst.scarlet.lite.Image`
90 ``convolved`` + ``noise``; the observation tests run against.
91 mCoadd : `lsst.afw.image.MultibandExposure`
92 ``noisy_image`` wrapped as the multiband exposure consumed by
93 detection and deblending.
94 psfs : `tuple` [`lsst.afw.detection.GaussianPsf`]
95 Per-band image-space PSF objects attached to ``mCoadd``.
96 modelPsf : `numpy.ndarray`
97 Pixelized narrow model PSF used to render ``deconvolved``.
98 imagePsf : `numpy.ndarray`
99 Pixelized per-band image PSFs used to render ``convolved``.
100 """
102 scene: Scene
103 deconvolved: scl.Image
104 convolved: scl.Image
105 noise: np.ndarray
106 noisy_image: scl.Image
107 mCoadd: afwImage.MultibandExposure
108 psfs: tuple[GaussianPsf, ...]
109 modelPsf: np.ndarray
110 imagePsf: np.ndarray
112 @property
113 def bands(self) -> tuple[str, ...]:
114 """Band ordering for this bundle (`tuple` [`str`], read-only)."""
115 return self.scene.bands
118@dataclass(frozen=True)
119class DetectionBundle:
120 """Detection stage of the pipeline for one (scene, config) pair.
122 Parameters
123 ----------
124 image : `ImageBundle`
125 Upstream rendered-image bundle that was detected on.
126 catalog : `lsst.afw.table.SourceCatalog`
127 Catalog of detected sources in the output (mapped) schema.
128 schema : `lsst.afw.table.Schema`
129 Output schema, after the schema mapper has been applied.
130 inputSchema : `lsst.afw.table.Schema`
131 Minimal input schema that the detection task was constructed
132 with.
133 schemaMapper : `lsst.afw.table.SchemaMapper`
134 Mapper from ``inputSchema`` to ``schema``.
135 detectionTask : `lsst.meas.algorithms.SourceDetectionTask`
136 The detection task instance, kept for inspection in tests.
137 detection_config_key : `str`
138 Serialized form of the detection config used as the cache key.
139 """
141 image: ImageBundle
142 catalog: SourceCatalog
143 schema: Schema
144 inputSchema: Schema
145 schemaMapper: SchemaMapper
146 detectionTask: SourceDetectionTask
147 detection_config_key: str
150@dataclass(frozen=True)
151class DeconvolveBundle:
152 """Deconvolution stage of the pipeline for one config tuple.
154 Parameters
155 ----------
156 detection : `DetectionBundle`
157 Upstream detection bundle that was deconvolved.
158 mDeconvolved : `lsst.afw.image.MultibandExposure`
159 Per-band deconvolved exposures produced by
160 `DeconvolveExposureTask`.
161 deconvolveTask : \
162 `lsst.meas.extensions.scarlet.DeconvolveExposureTask`
163 The deconvolve task instance, kept for inspection in tests.
164 deconvolve_config_key : `str`
165 Serialized form of the deconvolve config used as the cache key.
166 """
168 detection: DetectionBundle
169 mDeconvolved: afwImage.MultibandExposure
170 deconvolveTask: DeconvolveExposureTask
171 deconvolve_config_key: str
173 @property
174 def image(self) -> ImageBundle:
175 """Rendered-image bundle this stage was built from \
176(`ImageBundle`, read-only)."""
177 return self.detection.image
180@dataclass(frozen=True)
181class DeblendBundle:
182 """Deblend stage of the pipeline for one config tuple.
184 Parameters
185 ----------
186 deconvolved : `DeconvolveBundle`
187 Upstream deconvolved bundle that was deblended.
188 result : `lsst.pipe.base.Struct`
189 Raw result struct returned by ``ScarletDeblendTask.run``.
190 deblendTask : \
191 `lsst.meas.extensions.scarlet.ScarletDeblendTask`
192 The deblend task instance, kept for inspection in tests.
193 deblend_config_key : `str`
194 Serialized form of the deblend config used as the cache key.
195 """
197 deconvolved: DeconvolveBundle
198 result: Struct
199 deblendTask: ScarletDeblendTask
200 deblend_config_key: str
202 @property
203 def detection(self) -> DetectionBundle:
204 """Detection bundle this stage was built from \
205(`DetectionBundle`, read-only)."""
206 return self.deconvolved.detection
208 @property
209 def image(self) -> ImageBundle:
210 """Rendered-image bundle this stage was built from \
211(`ImageBundle`, read-only)."""
212 return self.detection.image
215def _config_key(config: Any) -> str:
216 """Serialize a task config to a stable cache key.
218 Parameters
219 ----------
220 config : `lsst.pex.config.Config` or `None`
221 Task config to key, or `None` to mean "use the task's default
222 config".
224 Returns
225 -------
226 key : `str`
227 ``"<default>"`` when ``config`` is `None`; otherwise a JSON
228 dump of ``config.toDict()`` with sorted keys.
229 """
230 if config is None:
231 return "<default>"
232 return json.dumps(config.toDict(), sort_keys=True, default=str)
235def _build_psfs() -> tuple[tuple[GaussianPsf, ...], np.ndarray, np.ndarray]:
236 """Build the per-band image PSFs and the narrow model PSF.
238 Returns
239 -------
240 psfs : `tuple` [`lsst.afw.detection.GaussianPsf`]
241 One `GaussianPsf` per entry in ``_PSF_SIGMAS``.
242 modelPsf : `numpy.ndarray`
243 Pixelized narrow model PSF, normalized.
244 imagePsf : `numpy.ndarray`
245 Pixelized per-band image PSFs, each normalized to unit sum.
246 """
247 modelPsf = scl.utils.integrated_circular_gaussian(
248 sigma=_MODEL_PSF_SIGMA
249 ).astype(np.float32)
250 psfShape = (2 * _PSF_RADIUS + 1, 2 * _PSF_RADIUS + 1)
251 psfs = tuple(
252 GaussianPsf(psfShape[1], psfShape[0], sigma) for sigma in _PSF_SIGMAS
253 )
254 imagePsf = np.asarray(
255 [psf.computeImage(psf.getAveragePosition()).array for psf in psfs]
256 ).astype(np.float32)
257 imagePsf /= imagePsf.sum(axis=(1, 2))[:, None, None]
258 return psfs, modelPsf, imagePsf
261# PSF objects built once at import time and shared across every scene.
262_PSFS, _MODEL_PSF, _IMAGE_PSF = _build_psfs()
265# Memoization caches, one per stage. Keys are tuples of all upstream
266# config keys plus this stage's own, so reusing a scene + config tuple
267# returns the same bundle without re-running the work.
268_image_cache: dict[str, ImageBundle] = {}
269_detection_cache: dict[tuple[str, str], DetectionBundle] = {}
270_deconvolve_cache: dict[tuple[str, str, str], DeconvolveBundle] = {}
271_deblend_cache: dict[tuple[str, str, str, str], DeblendBundle] = {}
274def build_image(scene: Scene) -> ImageBundle:
275 """Render the deconvolved, convolved, and noisy multiband image.
277 Parameters
278 ----------
279 scene : `Scene`
280 Scene whose models will be rendered.
282 Returns
283 -------
284 bundle : `ImageBundle`
285 Bundle holding both the truth (``deconvolved``) and the
286 observation (``noisy_image`` / ``mCoadd``), keyed on
287 ``scene.name``.
288 """
289 if scene.name in _image_cache:
290 return _image_cache[scene.name]
292 deconvolved, convolved = initData(scene.models, _MODEL_PSF, _IMAGE_PSF)
294 rng = np.random.RandomState(_NOISE_SEED)
295 noise = _NOISE_AMPLITUDE * (
296 rng.rand(*convolved.shape).astype(np.float32) - 0.5
297 )
298 noisy_image = convolved.copy()
299 noisy_image._data += noise
301 masked = afwImage.MultibandMaskedImage.fromArrays(
302 scene.bands, noisy_image.data, None, noise ** 2
303 )
304 coadds = [
305 afwImage.Exposure(img, dtype=img.image.array.dtype) for img in masked
306 ]
307 mCoadd = afwImage.MultibandExposure.fromExposures(scene.bands, coadds)
308 for b, coadd in enumerate(mCoadd):
309 coadd.setPsf(_PSFS[b])
311 bundle = ImageBundle(
312 scene=scene,
313 deconvolved=deconvolved,
314 convolved=convolved,
315 noise=noise,
316 noisy_image=noisy_image,
317 mCoadd=mCoadd,
318 psfs=_PSFS,
319 modelPsf=_MODEL_PSF,
320 imagePsf=_IMAGE_PSF,
321 )
322 _image_cache[scene.name] = bundle
323 return bundle
326def detect(image: ImageBundle, config: Any = None) -> DetectionBundle:
327 """Run :class:`SourceDetectionTask` on the r-band coadd.
329 Parameters
330 ----------
331 image : `ImageBundle`
332 Rendered-image bundle to detect on.
333 config : `lsst.meas.algorithms.SourceDetectionConfig`, optional
334 Detection-task config; the task's default is used when `None`.
336 Returns
337 -------
338 bundle : `DetectionBundle`
339 Bundle wrapping the catalog, schemas, and the task instance,
340 keyed on ``(image.scene.name, config_key)``.
341 """
342 config_key = _config_key(config)
343 key = (image.scene.name, config_key)
344 if key in _detection_cache:
345 return _detection_cache[key]
347 inputSchema = SourceTable.makeMinimalSchema()
348 table = SourceTable.make(inputSchema)
349 detectionTask = SourceDetectionTask(schema=inputSchema, config=config)
350 schemaMapper = SchemaMapper(inputSchema)
351 schemaMapper.addMinimalSchema(inputSchema)
352 schema = schemaMapper.getOutputSchema()
354 detectionResult = detectionTask.run(table, image.mCoadd["r"])
355 catalog_table = SourceCatalog.Table.make(schema)
356 catalog = SourceCatalog(catalog_table)
357 catalog.extend(detectionResult.sources, schemaMapper)
359 bundle = DetectionBundle(
360 image=image,
361 catalog=catalog,
362 schema=schema,
363 inputSchema=inputSchema,
364 schemaMapper=schemaMapper,
365 detectionTask=detectionTask,
366 detection_config_key=config_key,
367 )
368 _detection_cache[key] = bundle
369 return bundle
372def deconvolve(
373 detection: DetectionBundle, config: Any = None
374) -> DeconvolveBundle:
375 """Run :class:`DeconvolveExposureTask` once per band.
377 Parameters
378 ----------
379 detection : `DetectionBundle`
380 Detection-stage bundle whose coadds are to be deconvolved.
381 config : \
382 `lsst.meas.extensions.scarlet.DeconvolveExposureConfig`, \
383 optional
384 Deconvolve-task config; the task's default is used when `None`.
386 Returns
387 -------
388 bundle : `DeconvolveBundle`
389 Bundle wrapping the per-band deconvolved exposures and the
390 task instance, keyed on the upstream config keys plus this
391 stage's config key.
392 """
393 config_key = _config_key(config)
394 key = (
395 detection.image.scene.name,
396 detection.detection_config_key,
397 config_key,
398 )
399 if key in _deconvolve_cache:
400 return _deconvolve_cache[key]
402 deconvolveTask = DeconvolveExposureTask(config=config)
403 catalog = detection.catalog if deconvolveTask.config.useFootprints else None
405 deconvolvedCoadds = []
406 for coadd in detection.image.mCoadd:
407 deconvolvedCoadd = deconvolveTask.run(coadd, catalog).deconvolved
408 deconvolvedCoadds.append(deconvolvedCoadd)
409 mDeconvolved = afwImage.MultibandExposure.fromExposures(
410 detection.image.bands, deconvolvedCoadds
411 )
413 bundle = DeconvolveBundle(
414 detection=detection,
415 mDeconvolved=mDeconvolved,
416 deconvolveTask=deconvolveTask,
417 deconvolve_config_key=config_key,
418 )
419 _deconvolve_cache[key] = bundle
420 return bundle
423def deblend(
424 deconvolved: DeconvolveBundle, config: Any = None
425) -> DeblendBundle:
426 """Run :class:`ScarletDeblendTask`.
428 Parameters
429 ----------
430 deconvolved : `DeconvolveBundle`
431 Deconvolve-stage bundle whose coadds and catalog are
432 deblended.
433 config : \
434 `lsst.meas.extensions.scarlet.ScarletDeblendConfig`, \
435 optional
436 Deblend-task config; the task's default is used when `None`.
438 Returns
439 -------
440 bundle : `DeblendBundle`
441 Bundle wrapping the deblend result struct and the task
442 instance, keyed on the upstream config keys plus this stage's
443 config key.
444 """
445 config_key = _config_key(config)
446 key = (
447 deconvolved.image.scene.name,
448 deconvolved.detection.detection_config_key,
449 deconvolved.deconvolve_config_key,
450 config_key,
451 )
452 if key in _deblend_cache:
453 return _deblend_cache[key]
455 # ``ScarletDeblendTask.__init__`` adds deblend_* fields to the
456 # schema it receives. Clone the cached detection schema so a
457 # second call with a different config does not collide on the
458 # fields that the first call already added.
459 deblendTask = ScarletDeblendTask(
460 schema=Schema(deconvolved.detection.schema), config=config
461 )
462 result = deblendTask.run(
463 deconvolved.image.mCoadd,
464 deconvolved.mDeconvolved,
465 deconvolved.detection.catalog,
466 )
467 bundle = DeblendBundle(
468 deconvolved=deconvolved,
469 result=result,
470 deblendTask=deblendTask,
471 deblend_config_key=config_key,
472 )
473 _deblend_cache[key] = bundle
474 return bundle