Coverage for python/lsst/meas/extensions/scarlet/io/model_data.py: 94%
98 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-20 02:08 -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
24import json
25from typing import Any
27import numpy as np
28from numpy.typing import DTypeLike
30import lsst.scarlet.lite as scl
31from lsst.utils import DeprecatedDict
33from .hierarchical_blend_data import LEGACY_HIERARCHICAL_TYPES, LsstHierarchicalBlendData
34from .source_data import IsolatedSourceData
36__all__ = ["LsstScarletModelData"]
38CURRENT_SCHEMA = "1.0.2"
39MODEL_TYPE = "lsst"
40scl.io.migration.MigrationRegistry.set_current(MODEL_TYPE, CURRENT_SCHEMA)
43class LsstScarletModelData:
44 """A container that propagates scarlet models for an entire catalog,
45 including isolated sources.
47 This mirrors `~scarlet_lite.io.ScarletModelData` but carries information
48 specific to the LSST science pipelines, and owns its schema independent
49 from future changes to the scarlet_lite model.
51 Notes
52 -----
53 The persisted LsstScarletModelData object is stored to a zip file in the
54 science pipelines, where it's parameters are keys in the zip archive.
55 Those utilities are out of the migration registry scope, so we cheat a
56 little and package some of the attributes into a ``metadata`` dict
57 only for serialization. In :meth:`__init__` we lefit those fields
58 out of ``metadata`` and into typed attributes, and in :meth:`as_dict`
59 we fold them back into the metadata blob.
61 Attributes
62 ----------
63 blends
64 A mapping of parent IDs to blend data.
65 isolated
66 A mapping of isolated source IDs to their data.
67 metadata
68 A dictionary of additional metadata not needed for processing.
69 This is a `~lsst.utils.DeprecatedDict`: for a deprecation period
70 it also exposes ``bands``, ``model_psf`` and
71 ``psf`` as deprecated keys (mirrors of the typed attributes), which
72 warn on access and will be removed after v31.
73 bands
74 The ordered band labels of the model.
75 model_psf
76 The 2D model-space PSF shared by all bands.
77 psf
78 The per-band observed PSFs, shape ``(n_bands, height, width)``.
79 version
80 The schema version of the serialized data.
81 """
82 model_type: str = MODEL_TYPE
83 blends: dict[int, scl.io.ScarletBlendBaseData]
84 isolated: dict[int, IsolatedSourceData]
85 metadata: DeprecatedDict
86 version: str = CURRENT_SCHEMA
87 bands: tuple[str, ...] | None
88 model_psf: np.ndarray | None
89 psf: np.ndarray | None
91 def __init__(
92 self,
93 isolated: dict[int, IsolatedSourceData] | None = None,
94 blends: dict[int, scl.io.ScarletBlendBaseData] | None = None,
95 metadata: dict[str, Any] | None = None,
96 bands: tuple[str, ...] | None = None,
97 model_psf: np.ndarray | None = None,
98 psf: np.ndarray | None = None,
99 ):
100 self.blends = blends if blends is not None else {}
101 self.isolated = isolated if isolated is not None else {}
102 self.bands = bands
103 self.model_psf = model_psf
104 self.psf = psf
105 self.metadata = self._build_metadata(metadata, bands, model_psf, psf)
107 @staticmethod
108 def _build_metadata(
109 metadata: dict[str, Any] | None,
110 bands: tuple[str, ...] | None,
111 model_psf: np.ndarray | None,
112 psf: np.ndarray | None,
113 ) -> DeprecatedDict:
114 """Wrap ``metadata`` in a `DeprecatedDict`, injecting the promoted
115 typed attributes as deprecated back-compat keys.
116 """
117 data = dict(metadata) if metadata is not None else {}
118 if bands is not None:
119 data.setdefault("bands", tuple(bands))
120 if model_psf is not None:
121 data.setdefault("model_psf", model_psf)
122 if psf is not None:
123 data.setdefault("psf", psf)
124 return DeprecatedDict(
125 data,
126 deprecations={
127 key: (
128 f"Use the typed attribute LsstScarletModelData.{key} instead."
129 )
130 for key in ("bands", "model_psf", "psf")
131 },
132 version="v30.0",
133 )
135 def as_dict(self) -> dict[str, Any]:
136 """Convert to a dictionary for serialization
138 Returns
139 -------
140 result : dict[str, Any]
141 The object encoded as a JSON-compatible dictionary.
142 The mechanism for serializing to a zip file is outside of the
143 migration path, so the goal is to keep the result dict relatively
144 unchanged and hide new fields in the metadata blob, and extract
145 them in from_dict. So we should try to keep the result keys
146 as static as possible:
147 - ``model_type``: The type of the model, used for dispatch in
148 the migration registry.
149 - ``blends``: The dictionary of blend data.
150 - ``isolated``: The dictionary of isolated source data.
151 - ``metadata``: The metadata blob containing additional
152 information.
153 - ``version``: The schema version of the serialized data.
154 """
155 # Fold the typed attributes back into the metadata blob.
156 meta = dict(self.metadata) if self.metadata is not None else {}
157 if self.bands is not None: 157 ↛ 159line 157 didn't jump to line 159 because the condition on line 157 was always true
158 meta["bands"] = tuple(self.bands)
159 if self.model_psf is not None: 159 ↛ 161line 159 didn't jump to line 161 because the condition on line 159 was always true
160 meta["model_psf"] = self.model_psf
161 if self.psf is not None: 161 ↛ 163line 161 didn't jump to line 163 because the condition on line 161 was always true
162 meta["psf"] = self.psf
163 return {
164 "model_type": MODEL_TYPE,
165 "blends": {bid: b.as_dict() for bid, b in self.blends.items()},
166 "isolated": {sid: s.as_dict() for sid, s in self.isolated.items()},
167 "metadata": scl.io.utils.encode_metadata(meta) if meta else None,
168 "version": self.version,
169 }
171 def json(self) -> str:
172 """Serialize the data model to a JSON formatted string."""
173 return json.dumps(self.as_dict())
175 @classmethod
176 def from_dict(cls, data: dict, dtype: DTypeLike = np.float32) -> LsstScarletModelData:
177 """Reconstruct `LsstScarletModelData` from JSON compatible dict.
179 Parameters
180 ----------
181 data : dict
182 Dictionary representation of the object
183 dtype : DTypeLike
184 Datatype of the resulting model.
186 Returns
187 -------
188 result : LsstScarletModelData
189 The reconstructed object
190 """
191 data = scl.io.migration.MigrationRegistry.migrate(MODEL_TYPE, data)
192 blends: dict[int, scl.io.ScarletBlendBaseData] = {}
193 for bid, blend in data.get("blends", {}).items():
194 if "blend_type" not in blend:
195 # Default to a flat blend for legacy data.
196 blend["blend_type"] = "blend"
197 try:
198 blends[int(bid)] = scl.io.ScarletBlendBaseData.from_dict(blend, dtype=dtype)
199 except KeyError:
200 raise scl.io.utils.PersistenceError(
201 f"Unknown blend type: {blend['blend_type']} for blend ID: {bid}"
202 )
203 isolated: dict[int, IsolatedSourceData] = {}
204 for sid, source_data in data.get("isolated", {}).items():
205 isolated[int(sid)] = IsolatedSourceData.from_dict(source_data, dtype=dtype)
206 metadata = scl.io.utils.decode_metadata(data.get("metadata", None))
207 bands = metadata.pop("bands", None)
208 model_psf = metadata.pop("model_psf", None)
209 psf = metadata.pop("psf", None)
210 return cls(
211 isolated=isolated,
212 blends=blends,
213 metadata=metadata,
214 bands=bands,
215 model_psf=model_psf,
216 psf=psf
217 )
219 @classmethod
220 def parse_obj(cls, data: dict) -> LsstScarletModelData:
221 """Construct from a python-decoded JSON object (``json.load``)."""
222 return cls.from_dict(data, dtype=np.float32)
224 def to_scarlet_model_data(self) -> scl.io.ScarletModelData:
225 """Convert to a scarlet_lite `ScarletModelData` instance (lossy).
227 Returns
228 -------
229 result : `scarlet_lite.io.ScarletModelData`
230 A scarlet_lite model data instance with the shared blend and
231 metadata information.
232 """
233 return scl.io.ScarletModelData(blends=self.blends, metadata=self.metadata)
236@scl.io.migration.migration(MODEL_TYPE, scl.io.migration.PRE_SCHEMA)
237def _to_1_0_0(data: dict) -> dict:
238 """Migrate a pre-schema model to schema version 1.0.0
240 There were no changes to this data model in v1.0.0 but we need
241 to provide a way to migrate pre-schema data.
243 Parameters
244 ----------
245 data : dict
246 The data to migrate.
247 Returns
248 -------
249 result : dict
250 The migrated data.
251 """
252 # Ensure that the model type and version are set and add an
253 # empty isolated sources dictionary.
254 if "model_type" not in data: 254 ↛ 256line 254 didn't jump to line 256 because the condition on line 254 was always true
255 data["model_type"] = MODEL_TYPE
256 data["isolated"] = {}
257 # Pre-``metadata`` archives stored the model PSF as top-level
258 # ``psf`` / ``psfShape`` entries. Promote them into the modern
259 # ``metadata`` shape so ``decode_metadata`` can reconstruct the
260 # array via ``array_keys``. Mirrors scarlet_lite's pre-schema
261 # ``scarlet_model`` migration.
262 if "metadata" not in data and "psfShape" in data:
263 data["metadata"] = {
264 "model_psf": data.pop("psf"),
265 "model_psf_shape": data.pop("psfShape"),
266 "array_keys": ["model_psf"],
267 }
268 data["version"] = "1.0.0"
269 return data
272@scl.io.migration.migration(MODEL_TYPE, "1.0.0")
273def _to_1_0_1(data: dict) -> dict:
274 """Migrate a schema version 1.0.0 model to schema version 1.0.1
276 There were no changes to this data model in v1.0.1 but we need
277 to provide a way to migrate 1.0.0 data.
279 Parameters
280 ----------
281 data : dict
282 The data to migrate.
283 Returns
284 -------
285 result : dict
286 The migrated data.
287 """
288 data["version"] = "1.0.1"
289 if data.get("metadata") is None:
290 data["metadata"] = {}
291 data["metadata"].setdefault("footprint", None)
292 return data
295@scl.io.migration.migration(MODEL_TYPE, "1.0.1")
296def _to_1_0_2(data: dict) -> dict:
297 """Migrate a schema version 1.0.1 model to schema version 1.0.2.
299 1.0.1 (and earlier) stored top-level parent blends as scarlet_lite
300 ``HierarchicalBlendData`` with the detected-parent footprint
301 (``spans``/``origin``) in the blend ``metadata`` dict. 1.0.2 promotes each
302 to ``LsstHierarchicalBlendData``, where those fields are typed attributes.
304 Each legacy blend is handed to
305 ``LsstHierarchicalBlendData.convert_from_hierarchical``. This runs before
306 blend dispatch in ``from_dict``, so the rewritten dicts deserialize
307 directly to the new class.
309 Parameters
310 ----------
311 data : dict
312 The data to migrate.
314 Returns
315 -------
316 result : dict
317 The migrated data.
318 """
319 blends = data.get("blends", {})
320 for bid, blend in list(blends.items()):
321 if blend.get("blend_type", "blend") in LEGACY_HIERARCHICAL_TYPES:
322 blends[bid] = LsstHierarchicalBlendData.convert_from_hierarchical(blend)
323 data["version"] = "1.0.2"
324 return data