Coverage for python/lsst/meas/extensions/scarlet/io/source_data.py: 50%
54 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-08 01:48 -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 base64
25from dataclasses import dataclass
26from typing import Any
28import numpy as np
29from numpy.typing import DTypeLike
31import lsst.scarlet.lite as scl
32from ..source import IsolatedSource
34__all__ = ["IsolatedSourceData"]
36CURRENT_SCHEMA = "1.0.1"
37SOURCE_TYPE = "isolated"
38scl.io.migration.MigrationRegistry.set_current(SOURCE_TYPE, CURRENT_SCHEMA)
41def _encode_span_array(span_array: np.ndarray) -> str:
42 """Pack a 2D 0/1 mask into a base64-encoded packed-bits string.
44 The mask is flattened in C order, packed eight bits per byte via
45 :func:`numpy.packbits`, then base64-encoded. The companion
46 ``shape`` field carries the unflattened dimensions, since packing
47 pads to the next byte boundary.
49 Parameters
50 ----------
51 span_array : np.ndarray
52 Truthy values become 1; falsy become 0.
54 Returns
55 -------
56 encoded : str
57 ASCII base64 string suitable for direct JSON embedding.
58 """
59 flat = np.asarray(span_array, dtype=bool).ravel().astype(np.uint8)
60 return base64.b64encode(np.packbits(flat).tobytes()).decode("ascii")
63def _decode_span_array(
64 encoded: str, shape: tuple[int, ...], dtype: DTypeLike
65) -> np.ndarray:
66 """Decode the base64 packed-bits payload written by
67 :func:`_encode_span_array`.
69 Parameters
70 ----------
71 encoded : str
72 Base64 string from a 1.0.1+ ``span_array`` field.
73 shape : tuple[int, ...]
74 Output shape. ``np.prod(shape)`` defines how many bits are
75 valid; bytes past that are padding and are dropped.
76 dtype : DTypeLike
77 Output dtype. The decoded 0/1 values are cast to this type.
79 Returns
80 -------
81 span_array : np.ndarray
82 Mask of ``shape`` and ``dtype``.
83 """
84 packed = np.frombuffer(base64.b64decode(encoded), dtype=np.uint8)
85 n = int(np.prod(shape))
86 return np.unpackbits(packed)[:n].reshape(shape).astype(dtype)
89@dataclass(kw_only=True)
90class IsolatedSourceData(scl.io.blend.ScarletSourceBaseData):
91 """A source data instance of an isolated source.
93 This is used to represent sources that were not blended with any
94 other sources, and therefore do not have any deblending information.
96 Attributes
97 ----------
98 source_type : str
99 The type of source data.
100 version : str
101 The schema version of the serialized data.
102 span_array : np.ndarray
103 The span mask of the source.
104 origin : tuple[int, int]
105 The (y, x) origin of the footprint in the observation.
106 peak : tuple[int, int]
107 The (y, x) coordinates of the source peak in the observation.
108 metadata : dict | None
109 Additional metadata associated with the source.
110 """
112 source_type: str = SOURCE_TYPE
113 version: str = CURRENT_SCHEMA
114 span_array: np.ndarray
115 origin: tuple[int, int]
116 peak: tuple[int, int]
118 def as_dict(self) -> dict[str, Any]:
119 """Convert to a dictionary for serialization
121 Returns
122 -------
123 result : dict[str, Any]
124 The object encoded as a JSON-compatible dictionary.
125 """
126 result: dict[str, Any] = {
127 "origin": tuple(int(o) for o in self.origin),
128 "shape": tuple(int(s) for s in self.span_array.shape),
129 "peak": tuple(int(p) for p in self.peak),
130 "span_array": _encode_span_array(self.span_array),
131 "source_type": self.source_type,
132 "version": self.version,
133 }
134 if self.metadata is not None:
135 result["metadata"] = scl.io.utils.encode_metadata(self.metadata)
136 return result
138 @classmethod
139 def from_dict(cls, data: dict, dtype: DTypeLike = np.float32) -> IsolatedSourceData:
140 """Reconstruct `IsolatedSourceData` from JSON compatible
141 dict.
143 Parameters
144 ----------
145 data : dict
146 Dictionary representation of the object
147 dtype : DTypeLike
148 Datatype of the resulting model.
150 Returns
151 -------
152 result : IsolatedSourceData
153 The reconstructed object.
154 """
155 data = scl.io.migration.MigrationRegistry.migrate(SOURCE_TYPE, data)
156 shape = tuple(int(s) for s in data["shape"])
157 origin = tuple(int(o) for o in data["origin"])
158 span_array = _decode_span_array(data["span_array"], shape, dtype)
159 peak = tuple(int(p) for p in data["peak"])
160 metadata = scl.io.utils.decode_metadata(data.get("metadata", None))
161 return cls(
162 span_array=span_array,
163 origin=origin,
164 peak=peak,
165 metadata=metadata,
166 )
168 def to_source(self, observation: scl.Observation) -> IsolatedSource:
169 """Convert to a scarlet Source object
171 Parameters
172 ----------
173 observation : scl.Observation
174 The observation of the source.
176 Returns
177 -------
178 result : IsolatedSource
179 The scarlet Source object.
180 """
181 # Extract the image data that overlaps with the Footprint
182 bbox = scl.Box(self.span_array.shape, origin=self.origin)
183 image_data = observation.images[:, bbox].data
185 # Mask the image data with the footprint spans
186 model_data = image_data * self.span_array[None, :, :]
188 # Convert the array and bounding box into a scarlet Image
189 model = scl.Image(
190 model_data,
191 yx0=bbox.origin,
192 bands=observation.bands,
193 )
194 return IsolatedSource(model=model, peak=self.peak, metadata=self.metadata)
197IsolatedSourceData.register()
200@scl.io.migration.migration(SOURCE_TYPE, "1.0.0")
201def _to_1_0_1(data: dict) -> dict:
202 """Migrate an ``IsolatedSourceData`` payload from schema 1.0.0 to
203 1.0.1.
205 1.0.0 stored ``span_array`` as a flat tuple of Python floats
206 (one entry per pixel). 1.0.1 packs the same 0/1 mask into bytes
207 with :func:`numpy.packbits` and base64-encodes the result, which
208 cuts the on-disk representation by roughly an order of magnitude
209 for typical isolated-source footprints. The migration re-encodes
210 the legacy payload in place so the modern decoder handles it
211 transparently.
213 Parameters
214 ----------
215 data : dict
216 The data to migrate.
217 Returns
218 -------
219 result : dict
220 The migrated data.
221 """
222 shape = tuple(int(s) for s in data["shape"])
223 legacy = np.asarray(data["span_array"]).reshape(shape)
224 data["span_array"] = _encode_span_array(legacy)
225 data["version"] = "1.0.1"
226 return data