Coverage for python/lsst/meas/extensions/scarlet/io/utils.py: 16%
248 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-19 02:04 -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 if len(modelData.isolated) == 0:
180 return NoWorkFound("Scarlet model data is empty")
181 # All of the sources must have been isolated so there is nothing
182 # to do in this function. This is rare but it does occasionally
183 # happen in fields that only have u-band images.
184 return
185 if modelData.metadata is None:
186 raise ValueError("Scarlet model data does not contain metadata")
187 bands = modelData.metadata["bands"]
188 try:
189 bandIndex = bands.index(band)
190 except ValueError:
191 raise NoWorkFound(f"Band '{band}' not found in scarlet model data")
192 modelPsf = modelData.metadata["model_psf"]
193 observedPsf = modelData.metadata["psf"][bandIndex][None, :, :]
195 # Flux re-distribution may mix depth=1 blends, so we iterate over the
196 # completely flux separated parents to ensure that the full models
197 # are used for each source.
198 parents = catalog[catalog["parent"] == 0]
200 for n, parentRecord in enumerate(parents):
201 if parentRecord["deblend_nChild"] == 0:
202 # No children, so it is either an isolated source or failed
203 # deblending.
204 # Either way we can skip creating heavy child footprints.
205 continue
207 if updateFluxColumns and imageForRedistribution is not None:
208 # Update the data coverage (1 - # of NO_DATA pixels/# of pixels)
209 parentRecord["deblend_dataCoverage"] = calculateFootprintCoverage(
210 parentRecord.getFootprint(), imageForRedistribution.mask
211 )
213 observation = buildMonochromaticObservation(
214 modelPsf=modelPsf,
215 observedPsf=observedPsf,
216 footprint=parentRecord.getFootprint(),
217 imageForRedistribution=imageForRedistribution,
218 )
220 updateBlendRecords(
221 modelData=modelData,
222 bandIndex=bandIndex,
223 parent=parentRecord,
224 catalog=catalog,
225 observation=observation,
226 updateFluxColumns=updateFluxColumns,
227 imageForRedistribution=imageForRedistribution,
228 removeScarletData=removeScarletData,
229 )
232def buildMonochromaticObservation(
233 modelPsf: np.ndarray,
234 observedPsf: np.ndarray,
235 footprint: afwFootprint | None = None,
236 imageForRedistribution: MaskedImage | Exposure | None = None,
237) -> scl.Observation:
238 """Create a single-band observation for the entire image
240 Parameters
241 ----------
242 modelPsf :
243 The 2D model of the PSF.
244 observedPsf :
245 The observed PSF model for the catalog.
246 footprint :
247 The footprint of the source, used for masking out the model.
248 imageForRedistribution:
249 The image that is the source for flux re-distribution.
250 If `imageForRedistribution` is `None` then flux re-distribution is
251 not performed.
253 Returns
254 -------
255 observation : `scarlet.lite.Observation`
256 The observation for the entire image
257 """
258 if footprint is None and imageForRedistribution is None:
259 raise ValueError("Either footprint or imageForRedistribution must be provided")
261 if footprint is None:
262 bbox = imageForRedistribution.getBBox()
263 else:
264 bbox = footprint.getBBox()
266 scarletBox = utils.bboxToScarletBox(bbox)
268 if imageForRedistribution is not None:
269 cutout = imageForRedistribution[bbox]
271 # Mask the footprint
272 weights = np.ones(cutout.image.array.shape, dtype=cutout.image.array.dtype)
274 if footprint is not None:
275 weights *= footprint.spans.asArray()
277 observation = scl.Observation(
278 images=cutout.image.array[None, :, :],
279 variance=cutout.variance.array[None, :, :],
280 weights=weights[None, :, :],
281 psfs=observedPsf,
282 model_psf=modelPsf[None, :, :],
283 convolution_mode="real",
284 bands=monochromaticBands,
285 bbox=scarletBox,
286 )
287 else:
288 observation = scl.Observation.empty(
289 bands=monochromaticBands,
290 psfs=observedPsf,
291 model_psf=modelPsf[None, :, :],
292 bbox=scarletBox,
293 dtype=modelPsf.dtype,
294 )
295 return observation
298def calculateFootprintCoverage(footprint: afwFootprint, maskImage: MaskX) -> np.floating:
299 """Calculate the fraction of pixels with no data in a Footprint
301 Parameters
302 ----------
303 footprint : `lsst.afw.detection.Footprint`
304 The footprint to check for missing data.
305 maskImage : `lsst.afw.image.MaskX`
306 The mask image with the ``NO_DATA`` bit set.
307 Returns
308 -------
309 coverage : `float`
310 The fraction of pixels in `footprint` where the ``NO_DATA`` bit is set.
311 """
312 # Store the value of "NO_DATA" from the mask plane.
313 noDataInt = 2 ** maskImage.getMaskPlaneDict()["NO_DATA"]
315 # Calculate the coverage in the footprint
316 bbox = footprint.getBBox()
317 if bbox.area == 0:
318 # The source has no footprint, so it has no coverage
319 return 0
320 spans = footprint.spans.asArray()
321 totalArea = footprint.getArea()
322 mask = maskImage[bbox].array & noDataInt
323 noData = (mask * spans) > 0
324 coverage = 1 - np.sum(noData) / totalArea
325 return coverage
328def updateBlendRecords(
329 modelData: LsstScarletModelData,
330 bandIndex: int,
331 parent: SourceRecord,
332 catalog: SourceCatalog,
333 observation: scl.Observation,
334 updateFluxColumns: bool,
335 imageForRedistribution: MaskedImage | Exposure | None = None,
336 removeScarletData: bool = True,
337):
338 """Create footprints and update band-dependent columns in the catalog
340 Parameters
341 ----------
342 modelData :
343 Persistable data for the entire catalog.
344 bandIndex :
345 The number of the band to extract.
346 parent :
347 The parent source record.
348 catalog :
349 The catalog that is being updated.
350 observation :
351 The observation of the blend.
352 updateFluxColumns :
353 Whether or not to update the `deblend_*` columns in the catalog.
354 This should only be true when the input catalog schema already
355 contains those columns.
356 imageForRedistribution :
357 The image that is the source for flux re-distribution.
358 If `imageForRedistribution` is `None` then flux re-distribution is
359 not performed.
360 removeScarletData :
361 Whether or not to remove `ScarletBlendData` for each blend
362 to save memory.
363 """
364 useFlux = imageForRedistribution is not None
366 # Create a blend with the parent and all of its children.
367 blendData = modelData.blends[parent.getId()]
368 sources = []
369 if isinstance(blendData, scl.io.HierarchicalBlendData):
370 for blendId in blendData.children:
371 _blendData = cast(scl.io.ScarletBlendData, blendData.children[blendId])
372 blend = monochromaticDataToScarlet(_blendData, bandIndex, observation)
373 sources.extend(blend.sources)
375 if len(sources) == 0:
376 # No sources to update, so we can skip the rest of the function.
377 return
379 blend = scl.Blend(
380 sources=sources,
381 observation=observation,
382 )
384 if useFlux:
385 blend.conserve_flux()
387 # Set the metrics for the blend.
388 # TODO: remove this once DM-34558 runs all deblender metrics
389 # in a separate task.
390 if updateFluxColumns:
391 setDeblenderMetrics(blend)
393 # Update the HeavyFootprints for deblended sources
394 # and update the band-dependent catalog columns.
395 for source in blend.sources:
396 sourceRecord = catalog.find(source.metadata["id"])
398 peaks = parent.getFootprint().peaks
399 peakIdx = np.where(peaks["id"] == source.metadata["peak_id"])[0][0]
400 source.detectedPeak = peaks[peakIdx]
401 # Set the Footprint
402 heavy = scarletModelToHeavy(
403 source=source,
404 blend=blend,
405 useFlux=useFlux,
406 )
408 if updateFluxColumns:
409 if heavy.getArea() == 0:
410 # The source has no flux after being weighted with the PSF
411 # in this particular band (it might have flux in others).
412 sourceRecord.set("deblend_zeroFlux", True)
413 # Create a Footprint with a single pixel, set to zero,
414 # to avoid breakage in measurement algorithms.
415 center = Point2I(heavy.peaks[0]["i_x"], heavy.peaks[0]["i_y"])
416 spanList = [Span(center.y, center.x, center.x)]
417 footprint = afwFootprint(SpanSet(spanList))
418 footprint.setPeakCatalog(heavy.peaks)
419 heavy = HeavyFootprintF(footprint)
420 heavy.getImageArray()[0] = 0.0
421 else:
422 sourceRecord.set("deblend_zeroFlux", False)
423 sourceRecord.setFootprint(heavy)
425 if useFlux:
426 # Set the fraction of pixels with valid data.
427 coverage = calculateFootprintCoverage(
428 heavy, imageForRedistribution.mask
429 )
430 sourceRecord.set("deblend_dataCoverage", coverage)
432 # Set the flux of the scarlet model
433 # TODO: this field should probably be deprecated,
434 # since DM-33710 gives users access to the scarlet models.
435 model = source.get_model().data[0]
436 sourceRecord.set("deblend_scarletFlux", np.sum(model))
438 # Set the flux at the center of the model
439 peak = heavy.peaks[0]
441 img = heavy.extractImage(fill=0.0)
442 try:
443 sourceRecord.set(
444 "deblend_peak_instFlux", img[Point2I(peak["i_x"], peak["i_y"])]
445 )
446 except Exception:
447 srcId = sourceRecord.getId()
448 x = peak["i_x"]
449 y = peak["i_y"]
450 logger.warning(
451 f"Source {srcId} at {x},{y} could not set the peak flux with error:",
452 exc_info=1,
453 )
454 sourceRecord.set("deblend_peak_instFlux", np.nan)
455 else:
456 sourceRecord.setFootprint(heavy)
457 if removeScarletData:
458 if parent.getId() in modelData.blends:
459 del modelData.blends[parent.getId()]
462def build_scarlet_model(zip_dict: dict[str, Any]) -> LsstScarletModelData:
463 """Build a LsstScarletModelData instance from a dictionary of files.
465 Parameters
466 ----------
467 zip_dict : dict[str, Any]
468 Dictionary mapping filenames to the desired file type.
470 Returns
471 -------
472 model :
473 LsstScarletModelData instance.
474 """
475 metadata = zip_dict.pop('metadata', None)
476 version = zip_dict.pop('version', scl.io.migration.PRE_SCHEMA)
477 if metadata is None:
478 model_psf = zip_dict.pop('psf')
479 psf_shape = zip_dict.pop('psf_shape')
480 metadata = {
481 'psf': model_psf,
482 'psfShape': psf_shape,
483 }
484 blends = {}
485 isolated = {}
486 for key, value in zip_dict.items():
487 if "blend_type" in value:
488 blends[int(key)] = value
489 elif "source_type" in value:
490 if value["source_type"] != "isolated":
491 raise ValueError("Found unknown source type in scarlet model data isolated sources")
492 isolated[int(key)] = value
493 else:
494 raise ValueError(f"Found unknown file '{value}' in scarlet model data")
496 return LsstScarletModelData.parse_obj({
497 'version': version,
498 'isolated': isolated,
499 'blends': blends,
500 'metadata': metadata,
501 })
504def read_scarlet_model(path_or_stream: str, blend_ids: list[int] | None = None) -> LsstScarletModelData:
505 """Read a zip file and return a LsstScarletModelData instance.
507 Parameters
508 ----------
509 path : `str`
510 Path to the zip file.
511 blend_ids : `list[int]`, optional
512 List of blend IDs to extract from the zip file. If None,
513 all blends in the dataset will be extracted.
515 Returns
516 -------
517 model :
518 LsstScarletModelData instance.
519 """
521 if blend_ids is not None:
522 filenames = [str(f) for f in blend_ids]
523 else:
524 filenames = None
526 with zipfile.ZipFile(path_or_stream, 'r') as zip_file:
527 unzipped_files = {}
528 if filenames is None:
529 filenames = zip_file.namelist()
530 # Attempt to read the metadata file first, if it exists.
531 try:
532 with zip_file.open('metadata') as f:
533 metadata = from_json(f.read())
534 unzipped_files['metadata'] = metadata
535 except ValueError:
536 # The metadata file is not present, so we will
537 # assume that the model is in the legacy format.
538 filenames += ['psf', 'psf_shape']
539 try:
540 with zip_file.open('version') as f:
541 version = from_json(f.read())
542 unzipped_files['version'] = version
543 except KeyError:
544 # The version file is not present.
545 pass
547 for filename in filenames:
548 with zip_file.open(filename) as f:
549 unzipped_files[filename] = from_json(f.read())
551 return build_scarlet_model(unzipped_files)
554def scarlet_model_to_zip_json(model_data: LsstScarletModelData) -> dict[str, Any]:
555 """Convert a LsstScarletModelData instance to a dictionary of files.
557 This is required to convert the model data into a format that
558 can be insterted into a zip archive.
560 Parameters
561 ----------
562 model_data : `lsst.scarelt.lite.io.LsstScarletModelData`
563 LsstScarletModelData instance.
565 Returns
566 -------
567 data : dict[str, Any]
568 Dictionary mapping filenames to the desired file type.
569 """
570 json_model = model_data.as_dict()
572 data = {
573 str(blend_id): json.dumps(blend_data)
574 for blend_id, blend_data in json_model['blends'].items()
575 }
577 data.update({
578 str(source_id): json.dumps(source_data)
579 for source_id, source_data in json_model['isolated'].items()
580 })
581 # Support for legacy models
582 if 'psf' in json_model:
583 data.update({
584 'psf_shape': json.dumps(json_model['psfShape']),
585 'psf': json.dumps(json_model['psf']),
586 })
587 else:
588 data.update({
589 'metadata': json.dumps(json_model['metadata']),
590 'version': json.dumps(json_model['version']),
591 })
592 return data
595def write_scarlet_model(path_or_stream: str | BinaryIO, model_data: LsstScarletModelData):
596 """Write a LsstScarletModelData instance to a zip file.
598 Parameters
599 ----------
600 model_data : `lsst.scarlet.lite.io.LsstScarletModelData`
601 LsstScarletModelData instance.
603 Returns
604 -------
605 zip_dict :
606 Dictionary mapping filenames to the desired file type.
607 """
608 with zipfile.ZipFile(path_or_stream, 'w') as zf:
609 zip_archive = scarlet_model_to_zip_json(model_data)
610 for filename, data in zip_archive.items():
611 zf.writestr(filename, data)
614def scarlet_model_to_lsst_scarlet_model(model_data: scl.io.ScarletModelData) -> LsstScarletModelData:
615 """Convert a scarlet ModelData instance to a LsstScarletModelData instance.
617 Parameters
618 ----------
619 model_data : `scarlet.lite.io.ScarletModelData`
620 Scarlet ModelData instance.
622 Returns
623 -------
624 result : `lsst.scarlet.lite.io.LsstScarletModelData`
625 LsstScarletModelData instance.
626 """
627 return LsstScarletModelData(
628 blends=model_data.blends,
629 metadata=None,
630 )
633class ScarletModelFormatter(FormatterV2):
634 """Read and write scarlet models.
635 """
637 default_extension = ".scarlet"
638 unsupported_parameters = frozenset()
639 can_read_from_stream = True
640 can_read_from_local_file = True
642 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any:
643 # Override of `FormatterV2.read_from_local_file`.
644 return read_scarlet_model(path)
646 def read_from_stream(
647 self, stream: BinaryIO | ResourceHandleProtocol, component: str | None = None, expected_size: int = -1
648 ) -> Any:
649 # Override of `FormatterV2.read_from_stream`.
650 if self.file_descriptor.parameters is not None and "blend_id" in self.file_descriptor.parameters:
651 blend_ids = lsst_utils.iteration.ensure_iterable(self.file_descriptor.parameters["blend_id"])
652 else:
653 return NotImplemented
655 return read_scarlet_model(stream, blend_ids=blend_ids)
657 def to_bytes(self, in_memory_dataset: Any) -> bytes:
658 # Override of `FormatterV2.to_bytes`.
659 in_memory_zip = BytesIO()
660 write_scarlet_model(in_memory_zip, in_memory_dataset)
661 return in_memory_zip.getvalue()
664class ScarletModelDelegate(StorageClassDelegate):
665 """Delegate to extract a blend from an in-memory
666 LsstScarletModelData object.
667 """
668 def can_accept(self, inMemoryDataset: Any) -> bool:
669 return isinstance(inMemoryDataset, LsstScarletModelData)
671 def getComponent(self, composite: Any, componentName: str) -> Any:
672 raise AttributeError(f"Unsupported component: {componentName}")
674 def handleParameters(self, inMemoryDataset: Any, parameters: Mapping[str, Any] | None = None) -> Any:
675 if "blend_id" in parameters:
676 blend_ids = lsst_utils.iteration.ensure_iterable(parameters["blend_id"])
677 blends = {blend_id: inMemoryDataset.blends[blend_id] for blend_id in blend_ids}
678 inMemoryDataset.blends = blends
679 elif parameters is not None:
680 raise ValueError(f"Unsupported parameters: {parameters}")
681 return inMemoryDataset
684def loadBlend(blendData: scl.io.ScarletBlendData, model_psf: np.ndarray, mCoadd: MultibandExposure):
685 """Load a blend from the persisted data
687 Parameters
688 ----------
689 blendData:
690 The persisted scarlet BlendData to load into the blend.
691 model_psf:
692 The psf of the model in each band. This should be 2D, as scarlet
693 lite assumes that the PSF is the same for all bands.
694 mCoadd:
695 The coadd image to use for the observation attached to the blend.
696 This is required in order to create a difference kernel to convolve
697 the model into an observed seeing.
699 Returns
700 -------
701 blend : `scarlet.lite.Blend`
702 The blend object loaded from the persisted data.
703 """
704 psf, _ = utils.computePsfKernelImage(mCoadd, blendData.psf_center)
705 bbox = Box(blendData.shape, origin=blendData.origin)
706 afw_box = Box2I(Point2I(bbox.origin[::-1]), Extent2I(bbox.shape[::-1]))
707 coadd = mCoadd[blendData.bands, afw_box]
708 observation = scl.Observation(
709 images=coadd.image.array,
710 variance=coadd.variance.array,
711 weights=np.ones(coadd.image.array.shape, dtype=np.float32),
712 psfs=psf,
713 model_psf=model_psf[None, :, :],
714 convolution_mode='real',
715 bands=mCoadd.bands,
716 bbox=bbox,
717 )
718 return blendData.to_blend(observation), afw_box