Coverage for python/lsst/meas/extensions/scarlet/io/model_data.py: 94%
96 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-09 02:24 -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)
225@scl.io.migration.migration(MODEL_TYPE, scl.io.migration.PRE_SCHEMA)
226def _to_1_0_0(data: dict) -> dict:
227 """Migrate a pre-schema model to schema version 1.0.0
229 There were no changes to this data model in v1.0.0 but we need
230 to provide a way to migrate pre-schema data.
232 Parameters
233 ----------
234 data : dict
235 The data to migrate.
236 Returns
237 -------
238 result : dict
239 The migrated data.
240 """
241 # Ensure that the model type and version are set and add an
242 # empty isolated sources dictionary.
243 if "model_type" not in data: 243 ↛ 245line 243 didn't jump to line 245 because the condition on line 243 was always true
244 data["model_type"] = MODEL_TYPE
245 data["isolated"] = {}
246 # Pre-``metadata`` archives stored the model PSF as top-level
247 # ``psf`` / ``psfShape`` entries. Promote them into the modern
248 # ``metadata`` shape so ``decode_metadata`` can reconstruct the
249 # array via ``array_keys``. Mirrors scarlet_lite's pre-schema
250 # ``scarlet_model`` migration.
251 if "metadata" not in data and "psfShape" in data:
252 data["metadata"] = {
253 "model_psf": data.pop("psf"),
254 "model_psf_shape": data.pop("psfShape"),
255 "array_keys": ["model_psf"],
256 }
257 data["version"] = "1.0.0"
258 return data
261@scl.io.migration.migration(MODEL_TYPE, "1.0.0")
262def _to_1_0_1(data: dict) -> dict:
263 """Migrate a schema version 1.0.0 model to schema version 1.0.1
265 There were no changes to this data model in v1.0.1 but we need
266 to provide a way to migrate 1.0.0 data.
268 Parameters
269 ----------
270 data : dict
271 The data to migrate.
272 Returns
273 -------
274 result : dict
275 The migrated data.
276 """
277 data["version"] = "1.0.1"
278 if data.get("metadata") is None:
279 data["metadata"] = {}
280 data["metadata"].setdefault("footprint", None)
281 return data
284@scl.io.migration.migration(MODEL_TYPE, "1.0.1")
285def _to_1_0_2(data: dict) -> dict:
286 """Migrate a schema version 1.0.1 model to schema version 1.0.2.
288 1.0.1 (and earlier) stored top-level parent blends as scarlet_lite
289 ``HierarchicalBlendData`` with the detected-parent footprint
290 (``spans``/``origin``) in the blend ``metadata`` dict. 1.0.2 promotes each
291 to ``LsstHierarchicalBlendData``, where those fields are typed attributes.
293 Each legacy blend is handed to
294 ``LsstHierarchicalBlendData.convert_from_hierarchical``. This runs before
295 blend dispatch in ``from_dict``, so the rewritten dicts deserialize
296 directly to the new class.
298 Parameters
299 ----------
300 data : dict
301 The data to migrate.
303 Returns
304 -------
305 result : dict
306 The migrated data.
307 """
308 blends = data.get("blends", {})
309 for bid, blend in list(blends.items()):
310 if blend.get("blend_type", "blend") in LEGACY_HIERARCHICAL_TYPES:
311 blends[bid] = LsstHierarchicalBlendData.convert_from_hierarchical(blend)
312 data["version"] = "1.0.2"
313 return data