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