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

265 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-06-25 08:22 +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 .model_data import LsstScarletModelData 

58 

59logger = logging.getLogger(__name__) 

60 

61__all__ = [ 

62 "monochromaticDataToScarlet", 

63 "updateCatalogFootprints", 

64 "buildMonochromaticObservation", 

65 "calculateFootprintCoverage", 

66 "updateBlendRecords", 

67 "ScarletModelFormatter", 

68 "ScarletModelDelegate", 

69 "loadBlend", 

70] 

71 

72# Placeholder band label used by the deprecated 

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

74# from `modelData.metadata["bands"]`. Kept under a private name and 

75# surfaced under the deprecated public names via the module 

76# `__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.metadata['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.metadata is None: 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true

211 raise ValueError("Scarlet model data does not contain metadata") 

212 bands = modelData.metadata["bands"] 

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

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

215 bandIndex = bands.index(band) 

216 modelPsf = modelData.metadata["model_psf"] 

217 observedPsf = modelData.metadata["psf"][bandIndex][None, :, :] 

218 

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

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

221 # are used for each source. 

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

223 for parentId, blendData in blend_items: 

224 spans = blendData.metadata["spans"] 

225 bbox = scl.Box(spans.shape, blendData.metadata["origin"]) 

226 

227 observation = buildMonochromaticObservation( 

228 modelPsf=modelPsf, 

229 observedPsf=observedPsf, 

230 band=band, 

231 scarletBox=bbox, 

232 footprint=spans, 

233 imageForRedistribution=imageForRedistribution, 

234 ) 

235 

236 updateBlendRecords( 

237 modelData=modelData, 

238 blendData=blendData, 

239 band=band, 

240 catalog=catalog, 

241 observation=observation, 

242 updateFluxColumns=updateFluxColumns, 

243 imageForRedistribution=imageForRedistribution, 

244 ) 

245 

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

247 modelData.blends.pop(parentId, None) 

248 

249 

250def buildMonochromaticObservation( 

251 modelPsf: np.ndarray, 

252 observedPsf: np.ndarray, 

253 band: str, 

254 scarletBox: Box, 

255 footprint: np.ndarray | None, 

256 imageForRedistribution: MaskedImage | Exposure | None = None, 

257) -> scl.Observation: 

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

259 

260 Parameters 

261 ---------- 

262 modelPsf : 

263 The 2D model of the PSF. 

264 observedPsf : 

265 The observed PSF model for the catalog. 

266 band : 

267 Name of the band the observation represents. 

268 scarletBox : 

269 The bounding box for the scarlet observation. 

270 footprint : 

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

272 imageForRedistribution: 

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

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

275 not performed. 

276 

277 Returns 

278 ------- 

279 observation : `scarlet.lite.Observation` 

280 The observation for the entire image 

281 """ 

282 bbox = utils.scarletBoxToBBox(scarletBox) 

283 bands = (band,) 

284 

285 if imageForRedistribution is not None: 

286 cutout = imageForRedistribution[bbox] 

287 

288 # Mask the footprint 

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

290 

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

292 weights *= footprint 

293 

294 observation = scl.Observation( 

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

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

297 weights=weights[None, :, :], 

298 psfs=observedPsf, 

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

300 convolution_mode="real", 

301 bands=bands, 

302 bbox=scarletBox, 

303 ) 

304 else: 

305 observation = scl.Observation.empty( 

306 bands=bands, 

307 psfs=observedPsf, 

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

309 bbox=scarletBox, 

310 dtype=modelPsf.dtype, 

311 ) 

312 return observation 

313 

314 

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

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

317 

318 Parameters 

319 ---------- 

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

321 The footprint to check for missing data. 

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

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

324 

325 Returns 

326 ------- 

327 coverage : `float` 

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

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

330 ``deblend_dataCoverage`` catalog column. 

331 """ 

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

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

334 

335 # Calculate the coverage in the footprint 

336 bbox = footprint.getBBox() 

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

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

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

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

341 return np.float64(0.0) 

342 spans = footprint.spans.asArray() 

343 totalArea = footprint.getArea() 

344 mask = maskImage[bbox].array & noDataInt 

345 noData = (mask * spans) > 0 

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

347 return coverage 

348 

349 

350def updateBlendRecords( 

351 modelData: LsstScarletModelData, 

352 blendData: scl.io.ScarletBlendData | scl.io.HierarchicalBlendData, 

353 band: str, 

354 catalog: SourceCatalog, 

355 observation: scl.Observation, 

356 updateFluxColumns: bool, 

357 imageForRedistribution: MaskedImage | Exposure | None = None, 

358): 

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

360 

361 Parameters 

362 ---------- 

363 modelData : 

364 The full persisted scarlet model data. Its top-level 

365 ``metadata`` (``model_psf``, ``psf``, ``bands``) is used to 

366 reconstruct each child blend at full band-multiplicity before 

367 slicing to ``band``. 

368 blendData : 

369 Persistable data for a single blend or hierarchical blend. 

370 band : 

371 Name of the band to extract. 

372 catalog : 

373 The catalog that is being updated. 

374 observation : 

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

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

377 updateFluxColumns : 

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

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

380 contains those columns. 

381 imageForRedistribution : 

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

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

384 not performed. 

385 """ 

386 useFlux = imageForRedistribution is not None 

387 

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

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

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

391 sources = [] 

392 if isinstance(blendData, scl.io.HierarchicalBlendData): 392 ↛ 403line 392 didn't jump to line 403 because the condition on line 392 was always true

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

394 full_blend = cast( 

395 scl.io.ScarletBlendData, _blendData 

396 ).minimal_data_to_blend( 

397 model_psf=modelData.metadata["model_psf"][None, :, :], 

398 psf=modelData.metadata["psf"], 

399 bands=modelData.metadata["bands"], 

400 ) 

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

402 

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

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

405 return 

406 

407 blend = scl.Blend( 

408 sources=sources, 

409 observation=observation, 

410 ) 

411 

412 if useFlux: 

413 blend.conserve_flux() 

414 

415 # Set the metrics for the blend. 

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

417 setDeblenderMetrics(blend) 

418 

419 # Update the HeavyFootprints for deblended sources 

420 # and update the band-dependent catalog columns. 

421 for source in blend.sources: 

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

423 # Set the Footprint 

424 heavy = scarletModelToHeavy( 

425 source=source, 

426 blend=blend, 

427 useFlux=useFlux, 

428 ) 

429 

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

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

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

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

434 sourceRecord.set("deblend_zeroFlux", True) 

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

436 # to avoid breakage in measurement algorithms. 

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

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

439 footprint = afwFootprint(SpanSet(spanList)) 

440 footprint.setPeakCatalog(heavy.peaks) 

441 heavy = HeavyFootprintF(footprint) 

442 heavy.getImageArray()[0] = 0.0 

443 else: 

444 sourceRecord.set("deblend_zeroFlux", False) 

445 sourceRecord.setFootprint(heavy) 

446 

447 if useFlux: 

448 # Set the fraction of pixels with valid data. 

449 coverage = calculateFootprintCoverage( 

450 heavy, imageForRedistribution.mask 

451 ) 

452 sourceRecord.set("deblend_dataCoverage", coverage) 

453 

454 # Set the flux of the scarlet model 

455 # TODO: this field should probably be deprecated, 

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

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

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

459 

460 # Set the flux at the center of the model 

461 peak = heavy.peaks[0] 

462 

463 img = heavy.extractImage(fill=0.0) 

464 try: 

465 sourceRecord.set( 

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

467 ) 

468 except Exception: 

469 srcId = sourceRecord.getId() 

470 x = peak["i_x"] 

471 y = peak["i_y"] 

472 logger.warning( 

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

474 exc_info=True, 

475 ) 

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

477 

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

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

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

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

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

483 sourceRecord.set( 

484 "deblend_fluxOverlapFraction", metrics.fluxOverlapFraction[0] 

485 ) 

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

487 else: 

488 sourceRecord.setFootprint(heavy) 

489 

490 

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

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

493 

494 Parameters 

495 ---------- 

496 zip_dict : dict[str, Any] 

497 Dictionary mapping filenames to the desired file type. 

498 

499 Returns 

500 ------- 

501 model : 

502 LsstScarletModelData instance. 

503 """ 

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

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

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

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

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

509 # proper ``metadata`` dict. 

510 legacy_extras = {} 

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

512 if legacy_key in zip_dict: 

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

514 blends = {} 

515 isolated = {} 

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

517 if "blend_type" in value: 

518 blends[int(key)] = value 

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

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

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

522 isolated[int(key)] = value 

523 else: 

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

525 

526 payload: dict[str, Any] = { 

527 'version': version, 

528 'isolated': isolated, 

529 'blends': blends, 

530 } 

531 if metadata is not None: 

532 payload['metadata'] = metadata 

533 payload.update(legacy_extras) 

534 return LsstScarletModelData.parse_obj(payload) 

535 

536 

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

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

539 

540 Parameters 

541 ---------- 

542 path : `str` 

543 Path to the zip file. 

544 blend_ids : `list[int]`, optional 

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

546 all blends in the dataset will be extracted. 

547 

548 Returns 

549 ------- 

550 model : 

551 LsstScarletModelData instance. 

552 """ 

553 

554 if blend_ids is not None: 

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

556 else: 

557 filenames = None 

558 

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

560 unzipped_files = {} 

561 if filenames is None: 

562 filenames = zip_file.namelist() 

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

564 try: 

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

566 metadata = from_json(f.read()) 

567 unzipped_files['metadata'] = metadata 

568 except KeyError: 

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

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

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

572 try: 

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

574 version = from_json(f.read()) 

575 unzipped_files['version'] = version 

576 except KeyError: 

577 # The version file is not present. 

578 pass 

579 

580 for filename in filenames: 

581 with zip_file.open(filename) as f: 

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

583 

584 return build_scarlet_model(unzipped_files) 

585 

586 

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

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

589 

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

591 can be insterted into a zip archive. 

592 

593 Parameters 

594 ---------- 

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

596 LsstScarletModelData instance. 

597 

598 Returns 

599 ------- 

600 data : dict[str, Any] 

601 Dictionary mapping filenames to the desired file type. 

602 """ 

603 json_model = model_data.as_dict() 

604 

605 data = { 

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

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

608 } 

609 

610 data.update({ 

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

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

613 }) 

614 data.update({ 

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

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

617 }) 

618 return data 

619 

620 

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

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

623 

624 Parameters 

625 ---------- 

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

627 LsstScarletModelData instance. 

628 

629 Returns 

630 ------- 

631 zip_dict : 

632 Dictionary mapping filenames to the desired file type. 

633 """ 

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

635 zip_archive = scarlet_model_to_zip_json(model_data) 

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

637 zf.writestr(filename, data) 

638 

639 

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

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

642 

643 Parameters 

644 ---------- 

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

646 Scarlet ModelData instance. 

647 

648 Returns 

649 ------- 

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

651 LsstScarletModelData instance. 

652 """ 

653 return LsstScarletModelData( 

654 blends=model_data.blends, 

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

656 ) 

657 

658 

659class ScarletModelFormatter(FormatterV2): 

660 """Read and write scarlet models. 

661 """ 

662 

663 default_extension = ".scarlet" 

664 unsupported_parameters = frozenset() 

665 can_read_from_stream = True 

666 can_read_from_local_file = True 

667 

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

669 # Override of `FormatterV2.read_from_local_file`. 

670 return read_scarlet_model(path) 

671 

672 def read_from_stream( 

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

674 ) -> Any: 

675 # Override of `FormatterV2.read_from_stream`. 

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

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

678 else: 

679 return NotImplemented 

680 

681 return read_scarlet_model(stream, blend_ids=blend_ids) 

682 

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

684 # Override of `FormatterV2.to_bytes`. 

685 in_memory_zip = BytesIO() 

686 write_scarlet_model(in_memory_zip, in_memory_dataset) 

687 return in_memory_zip.getvalue() 

688 

689 

690class ScarletModelDelegate(StorageClassDelegate): 

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

692 LsstScarletModelData object. 

693 """ 

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

695 return isinstance(inMemoryDataset, LsstScarletModelData) 

696 

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

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

699 

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

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

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

703 # that here: the dispatch path in 

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

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

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

707 if not parameters: 

708 return inMemoryDataset 

709 if "blend_id" in parameters: 

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

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

712 inMemoryDataset.blends = blends 

713 else: 

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

715 return inMemoryDataset 

716 

717 

718def loadBlend( 

719 blendData: scl.io.ScarletBlendData, 

720 model_psf: np.ndarray | None = None, 

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

722 modelData: LsstScarletModelData | None = None, 

723): 

724 """Load a blend from the persisted data 

725 

726 Parameters 

727 ---------- 

728 blendData: 

729 The persisted scarlet BlendData to load into the blend. 

730 mCoadd: 

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

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

733 the model into an observed seeing. 

734 modelData: 

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

736 per-band ``psf`` and 2D ``model_psf`` stored in 

737 ``modelData.metadata`` are used to build the observation — 

738 these are the same PSFs the deblender saw during fitting and 

739 give a more faithful round-trip than re-deriving them from 

740 the coadd. 

741 model_psf: 

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

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

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

745 

746 Returns 

747 ------- 

748 blend : `scarlet.lite.Blend` 

749 The blend object loaded from the persisted data. 

750 afw_box : `lsst.geom.Box2I` 

751 The afw bounding box covering the blend. 

752 """ 

753 if model_psf is not None: 

754 warnings.warn( 

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

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

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

758 FutureWarning, stacklevel=2, 

759 ) 

760 if mCoadd is None: 

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

762 

763 psfs: np.ndarray 

764 if modelData is not None: 

765 if modelData.metadata is None: 765 ↛ 766line 765 didn't jump to line 766 because the condition on line 765 was never true

766 raise ValueError( 

767 "`modelData.metadata` must be populated to use the " 

768 "`modelData` branch of `loadBlend`." 

769 ) 

770 bands = tuple(modelData.metadata["bands"]) 

771 psfs = modelData.metadata["psf"] 

772 actual_model_psf = modelData.metadata["model_psf"][None, :, :] 

773 elif model_psf is not None: 

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

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

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

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

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

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

780 actual_model_psf = model_psf[None, :, :] 

781 else: 

782 raise ValueError( 

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

784 "to construct the observation." 

785 ) 

786 

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

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

789 coadd = mCoadd[bands, afw_box] 

790 observation = scl.Observation( 

791 images=coadd.image.array, 

792 variance=coadd.variance.array, 

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

794 psfs=psfs, 

795 model_psf=actual_model_psf, 

796 convolution_mode='real', 

797 bands=bands, 

798 bbox=bbox, 

799 ) 

800 return blendData.to_blend(observation), afw_box