Coverage for python/lsst/meas/extensions/scarlet/io/utils.py: 87%

268 statements  

« prev     ^ index     » next       coverage.py v7.15.0, created at 2026-07-09 09:32 +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/>. 

21 

22from __future__ import annotations 

23 

24from collections.abc import Mapping 

25from io import BytesIO 

26import logging 

27import json 

28from typing import Any, BinaryIO, cast 

29import warnings 

30import zipfile 

31 

32import numpy as np 

33from deprecated.sphinx import deprecated 

34from pydantic_core import from_json 

35 

36import lsst.scarlet.lite as scl 

37from lsst.afw.detection import Footprint as afwFootprint 

38from lsst.afw.detection import HeavyFootprintF 

39from lsst.afw.geom import Span, SpanSet 

40from lsst.afw.image import Exposure, MaskedImage, MaskX, MultibandExposure 

41from lsst.afw.table import SourceCatalog 

42import lsst.utils as lsst_utils 

43from lsst.daf.butler import StorageClassDelegate 

44from lsst.daf.butler import FormatterV2 

45from lsst.geom import Point2I, Extent2I, Box2I 

46from lsst.pipe.base import NoWorkFound 

47from lsst.resources import ResourceHandleProtocol 

48from lsst.scarlet.lite import ( 

49 Box, 

50 FactorizedComponent, 

51 FixedParameter, 

52) 

53 

54from ..metrics import setDeblenderMetrics 

55from .. import utils 

56from ..footprint import scarletModelToHeavy 

57from .hierarchical_blend_data import LsstHierarchicalBlendData 

58from .model_data import LsstScarletModelData 

59 

60logger = logging.getLogger(__name__) 

61 

62__all__ = [ 

63 "monochromaticDataToScarlet", 

64 "updateCatalogFootprints", 

65 "buildMonochromaticObservation", 

66 "calculateFootprintCoverage", 

67 "updateBlendRecords", 

68 "ScarletModelFormatter", 

69 "ScarletModelDelegate", 

70 "loadBlend", 

71] 

72 

73# Placeholder band label used by the deprecated 

74# `monochromaticDataToScarlet` flow. New code uses the real band name 

75# from `modelData.bands`. Kept under a private name and surfaced under 

76# the deprecated public names via the module `__getattr__` below. 

77_MONOCHROMATIC_BAND = "dummy" 

78_MONOCHROMATIC_BANDS = (_MONOCHROMATIC_BAND,) 

79 

80 

81def __getattr__(name: str) -> Any: 

82 if name in ("monochromaticBand", "monochromaticBands"): 

83 warnings.warn( 

84 f"`{name}` is deprecated and will be removed after v31; " 

85 "callers should use the real band name from " 

86 "`modelData.bands` instead.", 

87 FutureWarning, 

88 stacklevel=2, 

89 ) 

90 return _MONOCHROMATIC_BAND if name == "monochromaticBand" else _MONOCHROMATIC_BANDS 

91 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

92 

93 

94@deprecated( 

95 reason=( 

96 "Use `ScarletBlendData.minimal_data_to_blend(...)[band]` from " 

97 "`lsst.scarlet.lite` instead. Will be removed after v31." 

98 ), 

99 version="v30.0", 

100 category=FutureWarning, 

101) 

102def monochromaticDataToScarlet( 

103 blendData: scl.io.ScarletBlendData, 

104 bandIndex: int, 

105 observation: scl.Observation, 

106): 

107 """Convert the storage data model into a scarlet lite blend 

108 

109 Parameters 

110 ---------- 

111 blendData: 

112 Persistable data for the entire blend. 

113 bandIndex: 

114 Index of model to extract. 

115 observation: 

116 Observation of region inside the bounding box. 

117 

118 Returns 

119 ------- 

120 blend : `scarlet.lite.LiteBlend` 

121 A scarlet blend model extracted from persisted data. 

122 """ 

123 sources = [] 

124 # Use a dummy band, since we are only extracting a monochromatic model 

125 # that will be turned into a HeavyFootprint. 

126 bands = _MONOCHROMATIC_BANDS 

127 for sourceId, sourceData in blendData.sources.items(): 

128 components: list[scl.Component] = [] 

129 # There is no need to distinguish factorized components from regular 

130 # components, since there is only one band being used. 

131 for componentData in sourceData.components: 

132 if componentData.component_type == "component": 132 ↛ 133line 132 didn't jump to line 133 because the condition on line 132 was never true

133 bbox = Box(componentData.shape, origin=componentData.origin) 

134 model = scl.Image( 

135 componentData.model[bandIndex][None, :, :], yx0=bbox.origin, bands=bands 

136 ) 

137 component = scl.ComponentCube( 

138 model=model, 

139 peak=tuple(componentData.peak[::-1]), 

140 ) 

141 components.append(component) 

142 else: 

143 bbox = Box(componentData.shape, origin=componentData.origin) 

144 # Add dummy values for properties only needed for 

145 # model fitting. 

146 spectrum = FixedParameter(componentData.spectrum) 

147 totalBands = len(spectrum.x) 

148 morph = FixedParameter(componentData.morph) 

149 factorized = FactorizedComponent( 

150 bands=("dummy",) * totalBands, 

151 spectrum=spectrum, 

152 morph=morph, 

153 peak=tuple(int(np.round(p)) for p in componentData.peak), # type: ignore 

154 bbox=bbox, 

155 bg_rms=np.full((totalBands,), np.nan), 

156 ) 

157 model = factorized.get_model().data[bandIndex][None, :, :] 

158 model = scl.Image(model, yx0=bbox.origin, bands=bands) 

159 component = scl.component.CubeComponent( 

160 model=model, 

161 peak=factorized.peak, 

162 ) 

163 components.append(component) 

164 

165 source = scl.Source(components=components, metadata=sourceData.metadata) 

166 sources.append(source) 

167 

168 bbox = scl.Box(blendData.shape, origin=blendData.origin) 

169 blend = scl.Blend(sources=sources, observation=observation[:, bbox]) 

170 return blend 

171 

172 

173def updateCatalogFootprints( 

174 modelData: LsstScarletModelData, 

175 catalog: SourceCatalog, 

176 band: str, 

177 imageForRedistribution: MaskedImage | Exposure | None = None, 

178 removeScarletData: bool = True, 

179 updateFluxColumns: bool = True, 

180) -> None: 

181 """Use the scarlet models to set HeavyFootprints for modeled sources 

182 

183 Parameters 

184 ---------- 

185 modelData : 

186 Persistable data for the entire catalog. 

187 catalog: 

188 The catalog missing heavy footprints for deblended sources. 

189 band: 

190 The name of the band that the catalog data describes. 

191 imageForRedistribution: 

192 The image that is the source for flux re-distribution. 

193 If `imageForRedistribution` is `None` then flux re-distribution is 

194 not performed. 

195 removeScarletData: 

196 Whether or not to remove `ScarletBlendData` for each blend 

197 in order to save memory. 

198 updateFluxColumns: 

199 Whether or not to update the `deblend_*` columns in the catalog. 

200 This should only be true when the input catalog schema already 

201 contains those columns. 

202 """ 

203 if len(modelData.blends) == 0: 

204 if len(modelData.isolated) == 0: 

205 raise NoWorkFound("Scarlet model data is empty") 

206 # All of the sources must have been isolated so there is nothing 

207 # to do in this function. This is rare but it does occasionally 

208 # happen in fields that only have u-band images. 

209 return 

210 if modelData.bands is None or modelData.model_psf is None or modelData.psf is None: 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true

211 raise ValueError( 

212 "Scarlet model data is missing the `bands`, `model_psf` or `psf` " 

213 "needed to update catalog footprints." 

214 ) 

215 bands = modelData.bands 

216 if band not in bands: 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true

217 raise NoWorkFound(f"Band '{band}' not found in scarlet model data") 

218 bandIndex = bands.index(band) 

219 modelPsf = modelData.model_psf 

220 observedPsf = modelData.psf[bandIndex][None, :, :] 

221 

222 # Flux re-distribution may mix depth=1 blends, so we iterate over the 

223 # completely flux separated parents to ensure that the full models 

224 # are used for each source. 

225 blend_items = list(modelData.blends.items()) 

226 for parentId, blendData in blend_items: 

227 if not isinstance(blendData, LsstHierarchicalBlendData): 

228 raise ValueError( 

229 "updateCatalogFootprints expected an `LsstHierarchicalBlendData` " 

230 f"for parent {parentId}, got {type(blendData).__name__}. Top-level " 

231 "blends produced by the deblender are always hierarchical." 

232 ) 

233 spans = blendData.span_array 

234 bbox = scl.Box(spans.shape, blendData.origin) 

235 

236 observation = buildMonochromaticObservation( 

237 modelPsf=modelPsf, 

238 observedPsf=observedPsf, 

239 band=band, 

240 scarletBox=bbox, 

241 footprint=spans, 

242 imageForRedistribution=imageForRedistribution, 

243 ) 

244 

245 updateBlendRecords( 

246 modelData=modelData, 

247 blendData=blendData, 

248 band=band, 

249 catalog=catalog, 

250 observation=observation, 

251 updateFluxColumns=updateFluxColumns, 

252 imageForRedistribution=imageForRedistribution, 

253 ) 

254 

255 if removeScarletData: 255 ↛ 256line 255 didn't jump to line 256 because the condition on line 255 was never true

256 modelData.blends.pop(parentId, None) 

257 

258 

259def buildMonochromaticObservation( 

260 modelPsf: np.ndarray, 

261 observedPsf: np.ndarray, 

262 band: str, 

263 scarletBox: Box, 

264 footprint: np.ndarray | None, 

265 imageForRedistribution: MaskedImage | Exposure | None = None, 

266) -> scl.Observation: 

267 """Create a single-band observation for the entire image 

268 

269 Parameters 

270 ---------- 

271 modelPsf : 

272 The 2D model of the PSF. 

273 observedPsf : 

274 The observed PSF model for the catalog. 

275 band : 

276 Name of the band the observation represents. 

277 scarletBox : 

278 The bounding box for the scarlet observation. 

279 footprint : 

280 The footprint of the source, used for masking out the model. 

281 imageForRedistribution: 

282 The image that is the source for flux re-distribution. 

283 If `imageForRedistribution` is `None` then flux re-distribution is 

284 not performed. 

285 

286 Returns 

287 ------- 

288 observation : `scarlet.lite.Observation` 

289 The observation for the entire image 

290 """ 

291 bbox = utils.scarletBoxToBBox(scarletBox) 

292 bands = (band,) 

293 

294 if imageForRedistribution is not None: 

295 cutout = imageForRedistribution[bbox] 

296 

297 # Mask the footprint 

298 weights = np.ones(cutout.image.array.shape, dtype=cutout.image.array.dtype) 

299 

300 if footprint is not None: 300 ↛ 303line 300 didn't jump to line 303 because the condition on line 300 was always true

301 weights *= footprint 

302 

303 observation = scl.Observation( 

304 images=cutout.image.array[None, :, :], 

305 variance=cutout.variance.array[None, :, :], 

306 weights=weights[None, :, :], 

307 psfs=observedPsf, 

308 model_psf=modelPsf[None, :, :], 

309 convolution_mode="real", 

310 bands=bands, 

311 bbox=scarletBox, 

312 ) 

313 else: 

314 observation = scl.Observation.empty( 

315 bands=bands, 

316 psfs=observedPsf, 

317 model_psf=modelPsf[None, :, :], 

318 bbox=scarletBox, 

319 dtype=modelPsf.dtype, 

320 ) 

321 return observation 

322 

323 

324def calculateFootprintCoverage(footprint: afwFootprint, maskImage: MaskX) -> np.floating: 

325 """Calculate the fraction of pixels with valid data in a Footprint 

326 

327 Parameters 

328 ---------- 

329 footprint : `lsst.afw.detection.Footprint` 

330 The footprint to check for missing data. 

331 maskImage : `lsst.afw.image.MaskX` 

332 The mask image with the ``NO_DATA`` bit set. 

333 

334 Returns 

335 ------- 

336 coverage : `float` 

337 The fraction of pixels in `footprint` where the ``NO_DATA`` bit 

338 is **not** set (i.e. the fraction with valid data). Backs the 

339 ``deblend_dataCoverage`` catalog column. 

340 """ 

341 # Store the value of "NO_DATA" from the mask plane. 

342 noDataInt = 2 ** maskImage.getMaskPlaneDict()["NO_DATA"] 

343 

344 # Calculate the coverage in the footprint 

345 bbox = footprint.getBBox() 

346 if bbox.area == 0: 346 ↛ 350line 346 didn't jump to line 350 because the condition on line 346 was never true

347 # The source has no footprint, so it has no coverage. 

348 # Returning a ``np.float64`` (not a Python ``int``) honors the 

349 # ``-> np.floating`` annotation. 

350 return np.float64(0.0) 

351 spans = footprint.spans.asArray() 

352 totalArea = footprint.getArea() 

353 mask = maskImage[bbox].array & noDataInt 

354 noData = (mask * spans) > 0 

355 coverage = 1 - np.sum(noData) / totalArea 

356 return coverage 

357 

358 

359def updateBlendRecords( 

360 modelData: LsstScarletModelData, 

361 blendData: scl.io.ScarletBlendData | LsstHierarchicalBlendData, 

362 band: str, 

363 catalog: SourceCatalog, 

364 observation: scl.Observation, 

365 updateFluxColumns: bool, 

366 imageForRedistribution: MaskedImage | Exposure | None = None, 

367): 

368 """Create footprints and update band-dependent columns in the catalog 

369 

370 Parameters 

371 ---------- 

372 modelData : 

373 The full persisted scarlet model data. Its ``model_psf``, 

374 ``psf`` and ``bands`` attributes are used to reconstruct each 

375 child blend at full band-multiplicity before slicing to 

376 ``band``. 

377 blendData : 

378 Persistable data for a single blend or hierarchical blend. 

379 band : 

380 Name of the band to extract. 

381 catalog : 

382 The catalog that is being updated. 

383 observation : 

384 The observation of the blend in ``band``. Its ``bands`` tuple 

385 should be ``(band,)``. 

386 updateFluxColumns : 

387 Whether or not to update the `deblend_*` columns in the catalog. 

388 This should only be true when the input catalog schema already 

389 contains those columns. 

390 imageForRedistribution : 

391 The image that is the source for flux re-distribution. 

392 If `imageForRedistribution` is `None` then flux re-distribution is 

393 not performed. 

394 """ 

395 useFlux = imageForRedistribution is not None 

396 

397 # Reconstruct each child sub-blend from its persisted form, slice 

398 # down to the requested band, and collect the per-band sources into 

399 # a single Blend tied to the caller's redistribution observation. 

400 sources = [] 

401 if isinstance(blendData, LsstHierarchicalBlendData): 401 ↛ 412line 401 didn't jump to line 412 because the condition on line 401 was always true

402 for _blendData in blendData.children.values(): 

403 full_blend = cast( 

404 scl.io.ScarletBlendData, _blendData 

405 ).minimal_data_to_blend( 

406 model_psf=modelData.model_psf[None, :, :], 

407 psf=modelData.psf, 

408 bands=modelData.bands, 

409 ) 

410 sources.extend(full_blend[band].sources) 

411 

412 if len(sources) == 0: 412 ↛ 414line 412 didn't jump to line 414 because the condition on line 412 was never true

413 # No sources to update, so we can skip the rest of the function. 

414 return 

415 

416 blend = scl.Blend( 

417 sources=sources, 

418 observation=observation, 

419 ) 

420 

421 if useFlux: 

422 blend.conserve_flux() 

423 

424 # Set the metrics for the blend. 

425 if updateFluxColumns: 425 ↛ 430line 425 didn't jump to line 430 because the condition on line 425 was always true

426 setDeblenderMetrics(blend) 

427 

428 # Update the HeavyFootprints for deblended sources 

429 # and update the band-dependent catalog columns. 

430 for source in blend.sources: 

431 sourceRecord = catalog.find(source.metadata["id"]) 

432 # Set the Footprint 

433 heavy = scarletModelToHeavy( 

434 source=source, 

435 blend=blend, 

436 useFlux=useFlux, 

437 ) 

438 

439 if updateFluxColumns: 439 ↛ 497line 439 didn't jump to line 497 because the condition on line 439 was always true

440 if heavy.getArea() == 0: 440 ↛ 443line 440 didn't jump to line 443 because the condition on line 440 was never true

441 # The source has no flux after being weighted with the PSF 

442 # in this particular band (it might have flux in others). 

443 sourceRecord.set("deblend_zeroFlux", True) 

444 # Create a Footprint with a single pixel, set to zero, 

445 # to avoid breakage in measurement algorithms. 

446 center = Point2I(heavy.peaks[0]["i_x"], heavy.peaks[0]["i_y"]) 

447 spanList = [Span(center.y, center.x, center.x)] 

448 footprint = afwFootprint(SpanSet(spanList)) 

449 footprint.setPeakCatalog(heavy.peaks) 

450 heavy = HeavyFootprintF(footprint) 

451 heavy.getImageArray()[0] = 0.0 

452 else: 

453 sourceRecord.set("deblend_zeroFlux", False) 

454 sourceRecord.setFootprint(heavy) 

455 

456 if useFlux: 

457 # Set the fraction of pixels with valid data. 

458 coverage = calculateFootprintCoverage( 

459 heavy, imageForRedistribution.mask 

460 ) 

461 sourceRecord.set("deblend_dataCoverage", coverage) 

462 

463 # Set the flux of the scarlet model 

464 # TODO: this field should probably be deprecated, 

465 # since DM-33710 gives users access to the scarlet models. 

466 model = source.get_model().data[0] 

467 sourceRecord.set("deblend_scarletFlux", np.sum(model)) 

468 

469 # Set the flux at the center of the model 

470 peak = heavy.peaks[0] 

471 

472 img = heavy.extractImage(fill=0.0) 

473 try: 

474 sourceRecord.set( 

475 "deblend_peak_instFlux", img[Point2I(peak["i_x"], peak["i_y"])] 

476 ) 

477 except Exception: 

478 srcId = sourceRecord.getId() 

479 x = peak["i_x"] 

480 y = peak["i_y"] 

481 logger.warning( 

482 f"Source {srcId} at {x},{y} could not set the peak flux with error:", 

483 exc_info=True, 

484 ) 

485 sourceRecord.set("deblend_peak_instFlux", np.nan) 

486 

487 # The blend here is single-band, so every ``source.metrics`` 

488 # array has one entry — the value for ``band``. 

489 metrics = source.metrics # type: ignore[attr-defined] 

490 sourceRecord.set("deblend_maxOverlap", metrics.maxOverlap[0]) 

491 sourceRecord.set("deblend_fluxOverlap", metrics.fluxOverlap[0]) 

492 sourceRecord.set( 

493 "deblend_fluxOverlapFraction", metrics.fluxOverlapFraction[0] 

494 ) 

495 sourceRecord.set("deblend_blendedness", metrics.blendedness[0]) 

496 else: 

497 sourceRecord.setFootprint(heavy) 

498 

499 

500def build_scarlet_model(zip_dict: dict[str, Any]) -> LsstScarletModelData: 

501 """Build a LsstScarletModelData instance from a dictionary of files. 

502 

503 Parameters 

504 ---------- 

505 zip_dict : dict[str, Any] 

506 Dictionary mapping filenames to the desired file type. 

507 

508 Returns 

509 ------- 

510 model : 

511 LsstScarletModelData instance. 

512 """ 

513 metadata = zip_dict.pop('metadata', None) 

514 version = zip_dict.pop('version', scl.io.migration.PRE_SCHEMA) 

515 # Top-level legacy keys (e.g. ``psf`` / ``psfShape`` from a 

516 # pre-``metadata`` archive) are passed through to the migration 

517 # in ``LsstScarletModelData._to_1_0_0``, which synthesizes a 

518 # proper ``metadata`` dict. 

519 legacy_extras = {} 

520 for legacy_key in ('psf', 'psfShape'): 

521 if legacy_key in zip_dict: 

522 legacy_extras[legacy_key] = zip_dict.pop(legacy_key) 

523 blends = {} 

524 isolated = {} 

525 for key, value in zip_dict.items(): 

526 if "blend_type" in value: 

527 blends[int(key)] = value 

528 elif "source_type" in value: 528 ↛ 533line 528 didn't jump to line 533 because the condition on line 528 was always true

529 if value["source_type"] != "isolated": 529 ↛ 530line 529 didn't jump to line 530 because the condition on line 529 was never true

530 raise ValueError("Found unknown source type in scarlet model data isolated sources") 

531 isolated[int(key)] = value 

532 else: 

533 raise ValueError(f"Found unknown file '{value}' in scarlet model data") 

534 

535 payload: dict[str, Any] = { 

536 'version': version, 

537 'isolated': isolated, 

538 'blends': blends, 

539 } 

540 if metadata is not None: 

541 payload['metadata'] = metadata 

542 payload.update(legacy_extras) 

543 return LsstScarletModelData.parse_obj(payload) 

544 

545 

546def read_scarlet_model(path_or_stream: str, blend_ids: list[int] | None = None) -> LsstScarletModelData: 

547 """Read a zip file and return a LsstScarletModelData instance. 

548 

549 Parameters 

550 ---------- 

551 path : `str` 

552 Path to the zip file. 

553 blend_ids : `list[int]`, optional 

554 List of blend IDs to extract from the zip file. If None, 

555 all blends in the dataset will be extracted. 

556 

557 Returns 

558 ------- 

559 model : 

560 LsstScarletModelData instance. 

561 """ 

562 

563 if blend_ids is not None: 

564 filenames = [str(f) for f in blend_ids] 

565 else: 

566 filenames = None 

567 

568 with zipfile.ZipFile(path_or_stream, 'r') as zip_file: 

569 unzipped_files = {} 

570 if filenames is None: 

571 filenames = zip_file.namelist() 

572 # Attempt to read the metadata file first, if it exists. 

573 try: 

574 with zip_file.open('metadata') as f: 

575 metadata = from_json(f.read()) 

576 unzipped_files['metadata'] = metadata 

577 except KeyError: 

578 # The metadata file is not present, so we will 

579 # assume that the model is in the legacy format. 

580 filenames += ['psf', 'psfShape'] 

581 try: 

582 with zip_file.open('version') as f: 

583 version = from_json(f.read()) 

584 unzipped_files['version'] = version 

585 except KeyError: 

586 # The version file is not present. 

587 pass 

588 

589 for filename in filenames: 

590 with zip_file.open(filename) as f: 

591 unzipped_files[filename] = from_json(f.read()) 

592 

593 return build_scarlet_model(unzipped_files) 

594 

595 

596def scarlet_model_to_zip_json(model_data: LsstScarletModelData) -> dict[str, Any]: 

597 """Convert a LsstScarletModelData instance to a dictionary of files. 

598 

599 This is required to convert the model data into a format that 

600 can be insterted into a zip archive. 

601 

602 Parameters 

603 ---------- 

604 model_data : `lsst.meas.extensions.scarlet.io.LsstScarletModelData` 

605 LsstScarletModelData instance. 

606 

607 Returns 

608 ------- 

609 data : dict[str, Any] 

610 Dictionary mapping filenames to the desired file type. 

611 """ 

612 json_model = model_data.as_dict() 

613 

614 data = { 

615 str(blend_id): json.dumps(blend_data) 

616 for blend_id, blend_data in json_model['blends'].items() 

617 } 

618 

619 data.update({ 

620 str(source_id): json.dumps(source_data) 

621 for source_id, source_data in json_model['isolated'].items() 

622 }) 

623 data.update({ 

624 'metadata': json.dumps(json_model['metadata']), 

625 'version': json.dumps(json_model['version']), 

626 }) 

627 return data 

628 

629 

630def write_scarlet_model(path_or_stream: str | BinaryIO, model_data: LsstScarletModelData): 

631 """Write a LsstScarletModelData instance to a zip file. 

632 

633 Parameters 

634 ---------- 

635 model_data : `lsst.meas.extensions.scarlet.io.LsstScarletModelData` 

636 LsstScarletModelData instance. 

637 

638 Returns 

639 ------- 

640 zip_dict : 

641 Dictionary mapping filenames to the desired file type. 

642 """ 

643 with zipfile.ZipFile(path_or_stream, 'w') as zf: 

644 zip_archive = scarlet_model_to_zip_json(model_data) 

645 for filename, data in zip_archive.items(): 

646 zf.writestr(filename, data) 

647 

648 

649def scarlet_model_to_lsst_scarlet_model(model_data: scl.io.ScarletModelData) -> LsstScarletModelData: 

650 """Convert a scarlet ModelData instance to a LsstScarletModelData instance. 

651 

652 Parameters 

653 ---------- 

654 model_data : `scarlet.lite.io.ScarletModelData` 

655 Scarlet ModelData instance. 

656 

657 Returns 

658 ------- 

659 result : `lsst.meas.extensions.scarlet.io.LsstScarletModelData` 

660 LsstScarletModelData instance. 

661 """ 

662 return LsstScarletModelData( 

663 blends=model_data.blends, 

664 metadata=model_data.metadata if model_data.metadata is not None else {}, 

665 ) 

666 

667 

668class ScarletModelFormatter(FormatterV2): 

669 """Read and write scarlet models. 

670 """ 

671 

672 default_extension = ".scarlet" 

673 unsupported_parameters = frozenset() 

674 can_read_from_stream = True 

675 can_read_from_local_file = True 

676 

677 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any: 

678 # Override of `FormatterV2.read_from_local_file`. 

679 return read_scarlet_model(path) 

680 

681 def read_from_stream( 

682 self, stream: BinaryIO | ResourceHandleProtocol, component: str | None = None, expected_size: int = -1 

683 ) -> Any: 

684 # Override of `FormatterV2.read_from_stream`. 

685 if self.file_descriptor.parameters is not None and "blend_id" in self.file_descriptor.parameters: 685 ↛ 688line 685 didn't jump to line 688 because the condition on line 685 was always true

686 blend_ids = lsst_utils.iteration.ensure_iterable(self.file_descriptor.parameters["blend_id"]) 

687 else: 

688 return NotImplemented 

689 

690 return read_scarlet_model(stream, blend_ids=blend_ids) 

691 

692 def to_bytes(self, in_memory_dataset: Any) -> bytes: 

693 # Override of `FormatterV2.to_bytes`. 

694 in_memory_zip = BytesIO() 

695 write_scarlet_model(in_memory_zip, in_memory_dataset) 

696 return in_memory_zip.getvalue() 

697 

698 

699class ScarletModelDelegate(StorageClassDelegate): 

700 """Delegate to extract a blend from an in-memory 

701 LsstScarletModelData object. 

702 """ 

703 def can_accept(self, inMemoryDataset: Any) -> bool: 

704 return isinstance(inMemoryDataset, LsstScarletModelData) 

705 

706 def getComponent(self, composite: Any, componentName: str) -> Any: 

707 raise AttributeError(f"Unsupported component: {componentName}") 

708 

709 def handleParameters(self, inMemoryDataset: Any, parameters: Mapping[str, Any] | None = None) -> Any: 

710 # The base class signature permits ``parameters=None`` and 

711 # treats both ``None`` and ``{}`` as "no parameters". Mirror 

712 # that here: the dispatch path in 

713 # ``daf_butler.datastore.generic_base.post_process_get`` does 

714 # filter empty parameters before calling, but in-memory and 

715 # disassembled-composite reads pass ``None``/``{}`` through. 

716 if not parameters: 

717 return inMemoryDataset 

718 if "blend_id" in parameters: 

719 blend_ids = lsst_utils.iteration.ensure_iterable(parameters["blend_id"]) 

720 blends = {blend_id: inMemoryDataset.blends[blend_id] for blend_id in blend_ids} 

721 inMemoryDataset.blends = blends 

722 else: 

723 raise ValueError(f"Unsupported parameters: {parameters}") 

724 return inMemoryDataset 

725 

726 

727def loadBlend( 

728 blendData: scl.io.ScarletBlendData, 

729 model_psf: np.ndarray | None = None, 

730 mCoadd: MultibandExposure = None, # type: ignore[assignment] 

731 modelData: LsstScarletModelData | None = None, 

732): 

733 """Load a blend from the persisted data 

734 

735 Parameters 

736 ---------- 

737 blendData: 

738 The persisted scarlet BlendData to load into the blend. 

739 mCoadd: 

740 The coadd image to use for the observation attached to the blend. 

741 This is required in order to create a difference kernel to convolve 

742 the model into an observed seeing. 

743 modelData: 

744 The full persisted scarlet model data. When provided, the 

745 per-band ``psf`` and 2D ``model_psf`` carried on ``modelData`` 

746 are used to build the observation — these are the same PSFs the 

747 deblender saw during fitting and give a more faithful round-trip 

748 than re-deriving them from the coadd. 

749 model_psf: 

750 The 2D model-space PSF (deprecated). Retained for backward 

751 compatibility with the pre-``modelData`` signature; will be 

752 removed after v31. Passing it emits a ``FutureWarning``. 

753 

754 Returns 

755 ------- 

756 blend : `scarlet.lite.Blend` 

757 The blend object loaded from the persisted data. 

758 afw_box : `lsst.geom.Box2I` 

759 The afw bounding box covering the blend. 

760 """ 

761 if model_psf is not None: 

762 warnings.warn( 

763 "The `model_psf` parameter to `loadBlend` is deprecated and " 

764 "will be removed after v31; pass `modelData` instead so the " 

765 "PSFs from the fit can be reused directly.", 

766 FutureWarning, stacklevel=2, 

767 ) 

768 if mCoadd is None: 

769 raise ValueError("`mCoadd` is required to load a blend from persisted data") 

770 

771 psfs: np.ndarray 

772 if modelData is not None: 

773 if modelData.bands is None or modelData.psf is None or modelData.model_psf is None: 773 ↛ 774line 773 didn't jump to line 774 because the condition on line 773 was never true

774 raise ValueError( 

775 "`modelData` must carry `bands`, `psf` and `model_psf` to use " 

776 "the `modelData` branch of `loadBlend`." 

777 ) 

778 bands = tuple(modelData.bands) 

779 psfs = modelData.psf 

780 actual_model_psf = modelData.model_psf[None, :, :] 

781 elif model_psf is not None: 

782 # Legacy path: derive per-band PSFs from the coadd at the 

783 # blend's stored ``psf_center``. Modern ``ScarletBlendData`` 

784 # no longer carries ``psf_center`` or ``bands`` attributes, so 

785 # this branch only works against legacy-zip-read blends. 

786 psfs, _ = utils.computePsfKernelImage(mCoadd, blendData.psf_center) # type: ignore[attr-defined] 

787 bands = tuple(blendData.bands) # type: ignore[attr-defined] 

788 actual_model_psf = model_psf[None, :, :] 

789 else: 

790 raise ValueError( 

791 "loadBlend requires `modelData` (preferred) or `model_psf` " 

792 "to construct the observation." 

793 ) 

794 

795 bbox = Box(blendData.shape, origin=blendData.origin) 

796 afw_box = Box2I(Point2I(bbox.origin[::-1]), Extent2I(bbox.shape[::-1])) 

797 coadd = mCoadd[bands, afw_box] 

798 observation = scl.Observation( 

799 images=coadd.image.array, 

800 variance=coadd.variance.array, 

801 weights=np.ones(coadd.image.array.shape, dtype=np.float32), 

802 psfs=psfs, 

803 model_psf=actual_model_psf, 

804 convolution_mode='real', 

805 bands=bands, 

806 bbox=bbox, 

807 ) 

808 return blendData.to_blend(observation), afw_box