Coverage for python/lsst/meas/extensions/scarlet/io/utils.py: 18%
231 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-04 02:15 -0800
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/>.
22from __future__ import annotations
24from collections.abc import Mapping
25from io import BytesIO
26import logging
27import json
28from typing import Any, BinaryIO, cast
29import zipfile
31import numpy as np
32from pydantic_core import from_json
34import lsst.scarlet.lite as scl
35from lsst.afw.detection import Footprint as afwFootprint
36from lsst.afw.detection import HeavyFootprintF
37from lsst.afw.geom import Span, SpanSet
38from lsst.afw.image import Exposure, MaskedImage, MaskX, MultibandExposure
39from lsst.afw.table import SourceCatalog, SourceRecord
40import lsst.utils as lsst_utils
41from lsst.daf.butler import StorageClassDelegate
42from lsst.daf.butler import FormatterV2
43from lsst.geom import Point2I, Extent2I, Box2I
44from lsst.pipe.base import NoWorkFound
45from lsst.resources import ResourceHandleProtocol
46from lsst.scarlet.lite import (
47 Box,
48 FactorizedComponent,
49 FixedParameter,
50)
52from ..metrics import setDeblenderMetrics
53from .. import utils
54from ..footprint import scarletModelToHeavy
55from .model_data import LsstScarletModelData
57logger = logging.getLogger(__name__)
59__all__ = [
60 "monochromaticDataToScarlet",
61 "updateCatalogFootprints",
62 "buildMonochromaticObservation",
63 "calculateFootprintCoverage",
64 "updateBlendRecords",
65 "ScarletModelFormatter",
66 "ScarletModelDelegate",
67 "loadBlend",
68]
70# The name of the band in an monochome blend.
71# This is used as a placeholder since the band is not used in the
72# monochromatic model.
73monochromaticBand = "dummy"
74monochromaticBands = (monochromaticBand,)
77def monochromaticDataToScarlet(
78 blendData: scl.io.ScarletBlendData,
79 bandIndex: int,
80 observation: scl.Observation,
81):
82 """Convert the storage data model into a scarlet lite blend
84 Parameters
85 ----------
86 blendData:
87 Persistable data for the entire blend.
88 bandIndex:
89 Index of model to extract.
90 observation:
91 Observation of region inside the bounding box.
93 Returns
94 -------
95 blend : `scarlet.lite.LiteBlend`
96 A scarlet blend model extracted from persisted data.
97 """
98 sources = []
99 # Use a dummy band, since we are only extracting a monochromatic model
100 # that will be turned into a HeavyFootprint.
101 bands = monochromaticBands
102 for sourceId, sourceData in blendData.sources.items():
103 components: list[scl.Component] = []
104 # There is no need to distinguish factorized components from regular
105 # components, since there is only one band being used.
106 for componentData in sourceData.components:
107 if componentData.component_type == "component":
108 bbox = Box(componentData.shape, origin=componentData.origin)
109 model = scl.Image(
110 componentData.model[bandIndex][None, :, :], yx0=bbox.origin, bands=bands
111 )
112 component = scl.ComponentCube(
113 model=model,
114 peak=tuple(componentData.peak[::-1]),
115 )
116 components.append(component)
117 else:
118 bbox = Box(componentData.shape, origin=componentData.origin)
119 # Add dummy values for properties only needed for
120 # model fitting.
121 spectrum = FixedParameter(componentData.spectrum)
122 totalBands = len(spectrum.x)
123 morph = FixedParameter(componentData.morph)
124 factorized = FactorizedComponent(
125 bands=("dummy",) * totalBands,
126 spectrum=spectrum,
127 morph=morph,
128 peak=tuple(int(np.round(p)) for p in componentData.peak), # type: ignore
129 bbox=bbox,
130 bg_rms=np.full((totalBands,), np.nan),
131 )
132 model = factorized.get_model().data[bandIndex][None, :, :]
133 model = scl.Image(model, yx0=bbox.origin, bands=bands)
134 component = scl.io.ComponentCube(
135 model=model,
136 peak=factorized.peak,
137 )
138 components.append(component)
140 source = scl.Source(components=components, metadata=sourceData.metadata)
141 sources.append(source)
143 bbox = scl.Box(blendData.shape, origin=blendData.origin)
144 blend = scl.Blend(sources=sources, observation=observation[:, bbox])
145 return blend
148def updateCatalogFootprints(
149 modelData: LsstScarletModelData,
150 catalog: SourceCatalog,
151 band: str,
152 imageForRedistribution: MaskedImage | Exposure | None = None,
153 removeScarletData: bool = True,
154 updateFluxColumns: bool = True,
155) -> None:
156 """Use the scarlet models to set HeavyFootprints for modeled sources
158 Parameters
159 ----------
160 catalog:
161 The catalog missing heavy footprints for deblended sources.
162 band:
163 The name of the band that the catalog data describes.
164 imageForRedistribution:
165 The image that is the source for flux re-distribution.
166 If `imageForRedistribution` is `None` then flux re-distribution is
167 not performed.
168 removeScarletData:
169 Whether or not to remove `ScarletBlendData` for each blend
170 in order to save memory.
171 updateFluxColumns:
172 Whether or not to update the `deblend_*` columns in the catalog.
173 This should only be true when the input catalog schema already
174 contains those columns.
175 """
176 # All of the blends should have the same PSF,
177 # so we extract it from the first blend data.
178 if len(modelData.blends) == 0:
179 raise ValueError("Scarlet model data is empty")
180 if modelData.metadata is None:
181 raise ValueError("Scarlet model data does not contain metadata")
182 bands = modelData.metadata["bands"]
183 try:
184 bandIndex = bands.index(band)
185 except ValueError:
186 raise NoWorkFound(f"Band '{band}' not found in scarlet model data")
187 modelPsf = modelData.metadata["model_psf"]
188 observedPsf = modelData.metadata["psf"][bandIndex][None, :, :]
190 # Flux re-distribution may mix depth=1 blends, so we iterate over the
191 # completely flux separated parents to ensure that the full models
192 # are used for each source.
193 parents = catalog[catalog["parent"] == 0]
195 for n, parentRecord in enumerate(parents):
196 if parentRecord["deblend_nChild"] == 0:
197 # No children, so it is either an isolated source or failed
198 # deblending.
199 # Either way we can skip creating heavy child footprints.
200 continue
202 if updateFluxColumns and imageForRedistribution is not None:
203 # Update the data coverage (1 - # of NO_DATA pixels/# of pixels)
204 parentRecord["deblend_dataCoverage"] = calculateFootprintCoverage(
205 parentRecord.getFootprint(), imageForRedistribution.mask
206 )
208 observation = buildMonochromaticObservation(
209 modelPsf=modelPsf,
210 observedPsf=observedPsf,
211 footprint=parentRecord.getFootprint(),
212 imageForRedistribution=imageForRedistribution,
213 )
215 updateBlendRecords(
216 modelData=modelData,
217 bandIndex=bandIndex,
218 parent=parentRecord,
219 catalog=catalog,
220 observation=observation,
221 updateFluxColumns=updateFluxColumns,
222 imageForRedistribution=imageForRedistribution,
223 removeScarletData=removeScarletData,
224 )
227def buildMonochromaticObservation(
228 modelPsf: np.ndarray,
229 observedPsf: np.ndarray,
230 footprint: afwFootprint | None = None,
231 imageForRedistribution: MaskedImage | Exposure | None = None,
232) -> scl.Observation:
233 """Create a single-band observation for the entire image
235 Parameters
236 ----------
237 modelPsf :
238 The 2D model of the PSF.
239 observedPsf :
240 The observed PSF model for the catalog.
241 footprint :
242 The footprint of the source, used for masking out the model.
243 imageForRedistribution:
244 The image that is the source for flux re-distribution.
245 If `imageForRedistribution` is `None` then flux re-distribution is
246 not performed.
248 Returns
249 -------
250 observation : `scarlet.lite.Observation`
251 The observation for the entire image
252 """
253 if footprint is None and imageForRedistribution is None:
254 raise ValueError("Either footprint or imageForRedistribution must be provided")
256 if footprint is None:
257 bbox = imageForRedistribution.getBBox()
258 else:
259 bbox = footprint.getBBox()
261 scarletBox = utils.bboxToScarletBox(bbox)
263 if imageForRedistribution is not None:
264 cutout = imageForRedistribution[bbox]
266 # Mask the footprint
267 weights = np.ones(cutout.image.array.shape, dtype=cutout.image.array.dtype)
269 if footprint is not None:
270 weights *= footprint.spans.asArray()
272 observation = scl.Observation(
273 images=cutout.image.array[None, :, :],
274 variance=cutout.variance.array[None, :, :],
275 weights=weights[None, :, :],
276 psfs=observedPsf,
277 model_psf=modelPsf[None, :, :],
278 convolution_mode="real",
279 bands=monochromaticBands,
280 bbox=scarletBox,
281 )
282 else:
283 observation = scl.Observation.empty(
284 bands=monochromaticBands,
285 psfs=observedPsf,
286 model_psf=modelPsf[None, :, :],
287 bbox=scarletBox,
288 dtype=modelPsf.dtype,
289 )
290 return observation
293def calculateFootprintCoverage(footprint: afwFootprint, maskImage: MaskX) -> np.floating:
294 """Calculate the fraction of pixels with no data in a Footprint
296 Parameters
297 ----------
298 footprint : `lsst.afw.detection.Footprint`
299 The footprint to check for missing data.
300 maskImage : `lsst.afw.image.MaskX`
301 The mask image with the ``NO_DATA`` bit set.
302 Returns
303 -------
304 coverage : `float`
305 The fraction of pixels in `footprint` where the ``NO_DATA`` bit is set.
306 """
307 # Store the value of "NO_DATA" from the mask plane.
308 noDataInt = 2 ** maskImage.getMaskPlaneDict()["NO_DATA"]
310 # Calculate the coverage in the footprint
311 bbox = footprint.getBBox()
312 if bbox.area == 0:
313 # The source has no footprint, so it has no coverage
314 return 0
315 spans = footprint.spans.asArray()
316 totalArea = footprint.getArea()
317 mask = maskImage[bbox].array & noDataInt
318 noData = (mask * spans) > 0
319 coverage = 1 - np.sum(noData) / totalArea
320 return coverage
323def updateBlendRecords(
324 modelData: LsstScarletModelData,
325 bandIndex: int,
326 parent: SourceRecord,
327 catalog: SourceCatalog,
328 observation: scl.Observation,
329 updateFluxColumns: bool,
330 imageForRedistribution: MaskedImage | Exposure | None = None,
331 removeScarletData: bool = True,
332):
333 """Create footprints and update band-dependent columns in the catalog
335 Parameters
336 ----------
337 modelData :
338 Persistable data for the entire catalog.
339 bandIndex :
340 The number of the band to extract.
341 parent :
342 The parent source record.
343 catalog :
344 The catalog that is being updated.
345 observation :
346 The observation of the blend.
347 updateFluxColumns :
348 Whether or not to update the `deblend_*` columns in the catalog.
349 This should only be true when the input catalog schema already
350 contains those columns.
351 imageForRedistribution :
352 The image that is the source for flux re-distribution.
353 If `imageForRedistribution` is `None` then flux re-distribution is
354 not performed.
355 removeScarletData :
356 Whether or not to remove `ScarletBlendData` for each blend
357 to save memory.
358 """
359 useFlux = imageForRedistribution is not None
361 # Create a blend with the parent and all of its children.
362 blendData = modelData.blends[parent.getId()]
363 sources = []
364 if isinstance(blendData, scl.io.HierarchicalBlendData):
365 for blendId in blendData.children:
366 _blendData = cast(scl.io.ScarletBlendData, blendData.children[blendId])
367 blend = monochromaticDataToScarlet(_blendData, bandIndex, observation)
368 sources.extend(blend.sources)
370 if len(sources) == 0:
371 # No sources to update, so we can skip the rest of the function.
372 return
374 blend = scl.Blend(
375 sources=sources,
376 observation=observation,
377 )
379 if useFlux:
380 blend.conserve_flux()
382 # Set the metrics for the blend.
383 # TODO: remove this once DM-34558 runs all deblender metrics
384 # in a separate task.
385 if updateFluxColumns:
386 setDeblenderMetrics(blend)
388 # Update the HeavyFootprints for deblended sources
389 # and update the band-dependent catalog columns.
390 for source in blend.sources:
391 sourceRecord = catalog.find(source.metadata["id"])
393 peaks = parent.getFootprint().peaks
394 peakIdx = np.where(peaks["id"] == source.metadata["peak_id"])[0][0]
395 source.detectedPeak = peaks[peakIdx]
396 # Set the Footprint
397 heavy = scarletModelToHeavy(
398 source=source,
399 blend=blend,
400 useFlux=useFlux,
401 )
403 if updateFluxColumns:
404 if heavy.getArea() == 0:
405 # The source has no flux after being weighted with the PSF
406 # in this particular band (it might have flux in others).
407 sourceRecord.set("deblend_zeroFlux", True)
408 # Create a Footprint with a single pixel, set to zero,
409 # to avoid breakage in measurement algorithms.
410 center = Point2I(heavy.peaks[0]["i_x"], heavy.peaks[0]["i_y"])
411 spanList = [Span(center.y, center.x, center.x)]
412 footprint = afwFootprint(SpanSet(spanList))
413 footprint.setPeakCatalog(heavy.peaks)
414 heavy = HeavyFootprintF(footprint)
415 heavy.getImageArray()[0] = 0.0
416 else:
417 sourceRecord.set("deblend_zeroFlux", False)
418 sourceRecord.setFootprint(heavy)
420 if useFlux:
421 # Set the fraction of pixels with valid data.
422 coverage = calculateFootprintCoverage(
423 heavy, imageForRedistribution.mask
424 )
425 sourceRecord.set("deblend_dataCoverage", coverage)
427 # Set the flux of the scarlet model
428 # TODO: this field should probably be deprecated,
429 # since DM-33710 gives users access to the scarlet models.
430 model = source.get_model().data[0]
431 sourceRecord.set("deblend_scarletFlux", np.sum(model))
433 # Set the flux at the center of the model
434 peak = heavy.peaks[0]
436 img = heavy.extractImage(fill=0.0)
437 try:
438 sourceRecord.set(
439 "deblend_peak_instFlux", img[Point2I(peak["i_x"], peak["i_y"])]
440 )
441 except Exception:
442 srcId = sourceRecord.getId()
443 x = peak["i_x"]
444 y = peak["i_y"]
445 logger.warning(
446 f"Source {srcId} at {x},{y} could not set the peak flux with error:",
447 exc_info=1,
448 )
449 sourceRecord.set("deblend_peak_instFlux", np.nan)
450 else:
451 sourceRecord.setFootprint(heavy)
452 if removeScarletData:
453 if parent.getId() in modelData.blends:
454 del modelData.blends[parent.getId()]
457def build_scarlet_model(zip_dict: dict[str, Any]) -> LsstScarletModelData:
458 """Build a LsstScarletModelData instance from a dictionary of files.
460 Parameters
461 ----------
462 zip_dict : dict[str, Any]
463 Dictionary mapping filenames to the desired file type.
465 Returns
466 -------
467 model :
468 LsstScarletModelData instance.
469 """
470 metadata = zip_dict.pop('metadata', None)
471 if metadata is None:
472 model_psf = zip_dict.pop('psf')
473 psf_shape = zip_dict.pop('psf_shape')
474 metadata = {
475 'psf': model_psf,
476 'psfShape': psf_shape,
477 }
478 blends = {}
479 for key, value in zip_dict.items():
480 blends[int(key)] = value
481 return LsstScarletModelData.parse_obj({
482 'blends': blends,
483 'metadata': metadata,
484 })
487def read_scarlet_model(path_or_stream: str, blend_ids: list[int] | None = None) -> LsstScarletModelData:
488 """Read a zip file and return a LsstScarletModelData instance.
490 Parameters
491 ----------
492 path : `str`
493 Path to the zip file.
494 blend_ids : `list[int]`, optional
495 List of blend IDs to extract from the zip file. If None,
496 all blends in the dataset will be extracted.
498 Returns
499 -------
500 model :
501 LsstScarletModelData instance.
502 """
504 if blend_ids is not None:
505 filenames = [str(f) for f in blend_ids]
506 else:
507 filenames = None
509 with zipfile.ZipFile(path_or_stream, 'r') as zip_file:
510 unzipped_files = {}
511 if filenames is None:
512 filenames = zip_file.namelist()
513 # Attempt to read the metadata file first, if it exists.
514 try:
515 with zip_file.open('metadata') as f:
516 metadata = from_json(f.read())
517 unzipped_files['metadata'] = metadata
518 except ValueError:
519 # The metadata file is not present, so we will
520 # assume that the model is in the legacy format.
521 filenames += ['psf', 'psf_shape']
522 for filename in filenames:
523 with zip_file.open(filename) as f:
524 unzipped_files[filename] = from_json(f.read())
526 return build_scarlet_model(unzipped_files)
529def scarlet_model_to_zip_json(model_data: LsstScarletModelData) -> dict[str, Any]:
530 """Convert a LsstScarletModelData instance to a dictionary of files.
532 This is required to convert the model data into a format that
533 can be insterted into a zip archive.
535 Parameters
536 ----------
537 model_data : `lsst.scarelt.lite.io.LsstScarletModelData`
538 LsstScarletModelData instance.
540 Returns
541 -------
542 data : dict[str, Any]
543 Dictionary mapping filenames to the desired file type.
544 """
545 json_model = model_data.as_dict()
547 data = {
548 str(blend_id): json.dumps(blend_data)
549 for blend_id, blend_data in json_model['blends'].items()
550 }
551 # Support for legacy models
552 if 'psf' in json_model:
553 data.update({
554 'psf_shape': json.dumps(json_model['psfShape']),
555 'psf': json.dumps(json_model['psf']),
556 })
557 else:
558 data.update({
559 'metadata': json.dumps(json_model['metadata']),
560 })
561 return data
564def write_scarlet_model(path_or_stream: str | BinaryIO, model_data: LsstScarletModelData):
565 """Write a LsstScarletModelData instance to a zip file.
567 Parameters
568 ----------
569 model_data : `lsst.scarlet.lite.io.LsstScarletModelData`
570 LsstScarletModelData instance.
572 Returns
573 -------
574 zip_dict :
575 Dictionary mapping filenames to the desired file type.
576 """
577 with zipfile.ZipFile(path_or_stream, 'w') as zf:
578 zip_archive = scarlet_model_to_zip_json(model_data)
579 for filename, data in zip_archive.items():
580 zf.writestr(filename, data)
583def scarlet_model_to_lsst_scarlet_model(model_data: scl.io.ScarletModelData) -> LsstScarletModelData:
584 """Convert a scarlet ModelData instance to a LsstScarletModelData instance.
586 Parameters
587 ----------
588 model_data : `scarlet.lite.io.ScarletModelData`
589 Scarlet ModelData instance.
591 Returns
592 -------
593 result : `lsst.scarlet.lite.io.LsstScarletModelData`
594 LsstScarletModelData instance.
595 """
596 return LsstScarletModelData(
597 blends=model_data.blends,
598 metadata=None,
599 )
602class ScarletModelFormatter(FormatterV2):
603 """Read and write scarlet models.
604 """
606 default_extension = ".scarlet"
607 unsupported_parameters = frozenset()
608 can_read_from_stream = True
609 can_read_from_local_file = True
611 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any:
612 # Override of `FormatterV2.read_from_local_file`.
613 return read_scarlet_model(path)
615 def read_from_stream(
616 self, stream: BinaryIO | ResourceHandleProtocol, component: str | None = None, expected_size: int = -1
617 ) -> Any:
618 # Override of `FormatterV2.read_from_stream`.
619 if self.file_descriptor.parameters is not None and "blend_id" in self.file_descriptor.parameters:
620 blend_ids = lsst_utils.iteration.ensure_iterable(self.file_descriptor.parameters["blend_id"])
621 else:
622 return NotImplemented
624 return read_scarlet_model(stream, blend_ids=blend_ids)
626 def to_bytes(self, in_memory_dataset: Any) -> bytes:
627 # Override of `FormatterV2.to_bytes`.
628 in_memory_zip = BytesIO()
629 write_scarlet_model(in_memory_zip, in_memory_dataset)
630 return in_memory_zip.getvalue()
633class ScarletModelDelegate(StorageClassDelegate):
634 """Delegate to extract a blend from an in-memory
635 LsstScarletModelData object.
636 """
637 def can_accept(self, inMemoryDataset: Any) -> bool:
638 return isinstance(inMemoryDataset, LsstScarletModelData)
640 def getComponent(self, composite: Any, componentName: str) -> Any:
641 raise AttributeError(f"Unsupported component: {componentName}")
643 def handleParameters(self, inMemoryDataset: Any, parameters: Mapping[str, Any] | None = None) -> Any:
644 if "blend_id" in parameters:
645 blend_ids = lsst_utils.iteration.ensure_iterable(parameters["blend_id"])
646 blends = {blend_id: inMemoryDataset.blends[blend_id] for blend_id in blend_ids}
647 inMemoryDataset.blends = blends
648 elif parameters is not None:
649 raise ValueError(f"Unsupported parameters: {parameters}")
650 return inMemoryDataset
653def loadBlend(blendData: scl.io.ScarletBlendData, model_psf: np.ndarray, mCoadd: MultibandExposure):
654 """Load a blend from the persisted data
656 Parameters
657 ----------
658 blendData:
659 The persisted scarlet BlendData to load into the blend.
660 model_psf:
661 The psf of the model in each band. This should be 2D, as scarlet
662 lite assumes that the PSF is the same for all bands.
663 mCoadd:
664 The coadd image to use for the observation attached to the blend.
665 This is required in order to create a difference kernel to convolve
666 the model into an observed seeing.
668 Returns
669 -------
670 blend : `scarlet.lite.Blend`
671 The blend object loaded from the persisted data.
672 """
673 psf, _ = utils.computePsfKernelImage(mCoadd, blendData.psf_center)
674 bbox = Box(blendData.shape, origin=blendData.origin)
675 afw_box = Box2I(Point2I(bbox.origin[::-1]), Extent2I(bbox.shape[::-1]))
676 coadd = mCoadd[blendData.bands, afw_box]
677 observation = scl.Observation(
678 images=coadd.image.array,
679 variance=coadd.variance.array,
680 weights=np.ones(coadd.image.array.shape, dtype=np.float32),
681 psfs=psf,
682 model_psf=model_psf[None, :, :],
683 convolution_mode='real',
684 bands=mCoadd.bands,
685 bbox=bbox,
686 )
687 return blendData.to_blend(observation), afw_box