Coverage for tests/test_io_persistence.py: 98%
170 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-25 08:23 +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/>.
22"""Butler persistence tests for ``LsstScarletModelData``.
24Round-trips the deblender's on-disk model storage class plus two
25back-compatibility shims (a v1.0.0 ``LsstScarletModelData`` ingest and
26a v0 ``ScarletModelData`` ingest with a storage-class override). The
27deblend that supplies ``modelData`` comes from the cached pipeline
28stages in ``pipeline.py`` so that this file does not depend on
29``test_deblend.py``'s ad-hoc setup.
30"""
32import io
33import json
34import os
35import tempfile
36import unittest
37import zipfile
39import lsst.daf.butler
40import lsst.meas.extensions.scarlet as mes
41import lsst.scarlet.lite
42import lsst.utils.tests
43import numpy as np
44from lsst.daf.butler import (
45 Butler,
46 Config,
47 DatasetRef,
48 DatasetType,
49 FileDataset,
50 StorageClass,
51)
52from lsst.daf.butler.tests import makeTestCollection, makeTestRepo
54import pipeline
55from scenes import SCENES
57TESTDIR = os.path.abspath(os.path.dirname(__file__))
60class TestIoPersistence(lsst.utils.tests.TestCase):
61 """Butler put/get and legacy-model tests for
62 ``LsstScarletModelData`` storage in
63 ``lsst.meas.extensions.scarlet.io``.
64 """
66 def _persist_modelData(self):
67 # Set up a butler with the multi-blend modelData written into
68 # it. Sets ``self.modelData``, ``self.model_psf``, ``self.psf``,
69 # ``self.bands``, and ``self.butler`` for use by the three
70 # put/get tests below. ``pipeline.deblend`` is memoized per
71 # scene + config, so the deblend itself is computed once per
72 # process even though this helper runs per test.
73 bundle = pipeline.deblend(
74 pipeline.deconvolve(
75 pipeline.detect(pipeline.build_image(SCENES["multi-blend"]))
76 )
77 )
78 self.modelData = bundle.result.scarletModelData
79 self.bands = self.modelData.metadata["bands"]
80 self.model_psf = self.modelData.metadata["model_psf"][None, :, :]
81 self.psf = self.modelData.metadata["psf"]
82 repo = self._setup_butler()
83 self.butler = makeTestCollection(repo, uniqueId="test_run1")
84 self.butler.put(self.modelData, "scarlet_model_data", dataId={})
86 def test_butler_put_get_roundtrip(self):
87 """A butler ``put`` then ``get`` (no parameters) preserves
88 the full ``LsstScarletModelData``.
90 Checks ``model_psf`` and ``psf`` metadata, the blend count
91 and per-blend children (compared via ``_test_blend``), and the
92 isolated-source origins and span arrays.
93 """
94 self._persist_modelData()
95 modelData2 = self.butler.get("scarlet_model_data", dataId={})
97 np.testing.assert_almost_equal(
98 modelData2.metadata["model_psf"][None, :, :], self.model_psf
99 )
100 np.testing.assert_almost_equal(modelData2.metadata["psf"], self.psf)
101 self.assertEqual(len(modelData2.blends), len(self.modelData.blends))
103 for parentId in self.modelData.blends.keys():
104 nChildren = len(self.modelData.blends[parentId].children)
105 self.assertEqual(nChildren, len(modelData2.blends[parentId].children))
106 for blendId in self.modelData.blends[parentId].children:
107 blendData1 = self.modelData.blends[parentId].children[blendId]
108 blendData2 = modelData2.blends[parentId].children[blendId]
109 self._test_blend(blendData1, blendData2, self.model_psf, self.psf, self.bands)
111 for sourceId in self.modelData.isolated.keys():
112 isolatedData1 = self.modelData.isolated[sourceId]
113 isolatedData2 = modelData2.isolated[sourceId]
114 self.assertTupleEqual(isolatedData1.origin, isolatedData2.origin)
115 np.testing.assert_array_equal(
116 isolatedData1.span_array,
117 isolatedData2.span_array,
118 )
120 def test_butler_get_single_blend_parameter(self):
121 """``parameters={'blend_id': id}`` returns exactly that one blend.
123 The returned modelData contains only the requested parent and
124 its children are bit-identical (via ``_test_blend``) to the
125 original.
126 """
127 self._persist_modelData()
128 parentId = next(iter(self.modelData.blends))
130 modelData2 = self.butler.get(
131 "scarlet_model_data", dataId={}, parameters={"blend_id": parentId}
132 )
134 self.assertEqual(len(modelData2.blends), 1)
135 self.assertIn(parentId, modelData2.blends)
136 for blendId, blendData1 in self.modelData.blends[parentId].children.items():
137 blendData2 = modelData2.blends[parentId].children[blendId]
138 self._test_blend(blendData1, blendData2, self.model_psf, self.psf, self.bands)
140 def test_butler_get_multiple_blend_parameter(self):
141 """``parameters={'blend_id': [...]}`` returns exactly the listed
142 blends.
144 Picks the first two parent IDs from the multi-blend scene so the
145 test does not hardcode specific catalog IDs (which depend on
146 detection ordering).
147 """
148 self._persist_modelData()
149 blendIds = list(self.modelData.blends.keys())[:2]
151 modelData2 = self.butler.get(
152 "scarlet_model_data", dataId={}, parameters={"blend_id": blendIds}
153 )
155 self.assertEqual(len(modelData2.blends), len(blendIds))
156 for parentId in blendIds:
157 parentData1 = self.modelData.blends[parentId]
158 parentData2 = modelData2.blends[parentId]
159 self.assertEqual(len(parentData1.children), len(parentData2.children))
160 for blendId in parentData1.children.keys():
161 blendData1 = parentData1.children[blendId]
162 blendData2 = parentData2.children[blendId]
163 self._test_blend(blendData1, blendData2, self.model_psf, self.psf, self.bands)
165 def test_legacy_model(self):
166 repo = self._setup_butler()
167 storageClass = StorageClass(
168 "LsstScarletModelData",
169 pytype=mes.io.LsstScarletModelData,
170 )
171 datasetType = DatasetType(
172 "old_scarlet_model_data",
173 dimensions=(),
174 storageClass=storageClass,
175 universe=repo.dimensions,
176 )
177 ref = DatasetRef(
178 datasetType,
179 run="test_ingestion",
180 dataId={},
181 )
182 dataset = FileDataset(
183 path=os.path.join(TESTDIR, "data", "v29_models.json"),
184 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
185 refs=[ref],
186 )
188 # Ingest the legacy model into the butler
189 butler = makeTestCollection(repo, uniqueId="ingestion")
190 repo.registry.registerDatasetType(datasetType)
191 butler.ingest(dataset)
193 model = butler.get("old_scarlet_model_data", dataId={})
194 self.assertEqual(len(model.blends), 2)
195 # The pre-``metadata`` archive stored the model PSF as the
196 # top-level ``psf`` / ``psfShape`` entries. The legacy
197 # migration must promote those into ``metadata['model_psf']``
198 # (numpy array, reconstructed via ``array_keys``) so
199 # downstream consumers see the same shape as a modern model.
200 # Regression test for finding IO-17 of
201 # ``audits/audit-2026-05-05.md``.
202 self.assertIsNotNone(model.metadata)
203 self.assertIn("model_psf", model.metadata)
204 self.assertIsInstance(model.metadata["model_psf"], np.ndarray)
205 self.assertEqual(model.metadata["model_psf"].shape, (15, 15))
206 self.assertNotIn("psfShape", model.metadata)
208 test = butler.get("old_scarlet_model_data", dataId={}, parameters={"blend_id": 3495976385350991873})
209 self.assertEqual(len(test.blends), 1)
211 def test_v30_legacy_model(self):
212 """``LsstScarletModelData`` ingested from a v30-era fixture
213 round-trips intact.
215 ``data/v30_models.json`` snapshots the on-disk layout produced
216 by LSST release v30 — model schema ``1.0.1`` and
217 isolated-source schema ``1.0.0``. It plays the same role for
218 future schema bumps that ``v29_models.json`` plays for the
219 pre-``metadata`` layout: as long as the migration chain stays
220 complete, a v30 archive must continue to load against whatever
221 schemas a later release ships.
222 """
223 repo = self._setup_butler()
224 storageClass = StorageClass(
225 "LsstScarletModelData",
226 pytype=mes.io.LsstScarletModelData,
227 )
228 datasetType = DatasetType(
229 "old_scarlet_model_data",
230 dimensions=(),
231 storageClass=storageClass,
232 universe=repo.dimensions,
233 )
234 ref = DatasetRef(
235 datasetType,
236 run="test_ingestion_v30",
237 dataId={},
238 )
239 dataset = FileDataset(
240 path=os.path.join(TESTDIR, "data", "v30_models.json"),
241 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
242 refs=[ref],
243 )
245 butler = makeTestCollection(repo, uniqueId="ingestion_v30")
246 repo.registry.registerDatasetType(datasetType)
247 butler.ingest(dataset)
249 model = butler.get("old_scarlet_model_data", dataId={})
251 # The multi-blend scene that generated the fixture produces
252 # three parent blends and one isolated source.
253 self.assertEqual(len(model.blends), 3)
254 self.assertEqual(len(model.isolated), 1)
256 # Metadata round-trips with the model_psf array reconstructed
257 # via ``decode_metadata``'s ``array_keys`` handling.
258 self.assertIsNotNone(model.metadata)
259 self.assertIn("model_psf", model.metadata)
260 self.assertIsInstance(model.metadata["model_psf"], np.ndarray)
261 self.assertEqual(model.metadata["model_psf"].shape, (15, 15))
262 self.assertIn("psf", model.metadata)
263 self.assertIn("bands", model.metadata)
265 # The isolated source survives the full ``IsolatedSourceData``
266 # round-trip: shape and integer peak (post-IO-1), and a
267 # bit-exact span mask. Pinning the span sum guards the
268 # ``span_array`` serialization path against silent regressions
269 # under future schema bumps.
270 iso = next(iter(model.isolated.values()))
271 self.assertEqual(iso.span_array.shape, (13, 13))
272 self.assertEqual(iso.origin, (6, 14))
273 self.assertEqual(iso.peak, (12, 20))
274 self.assertEqual(float(iso.span_array.sum()), 119.0)
276 # Single-blend parameter load also works on v30 archives.
277 first_blend_id = sorted(model.blends.keys())[0]
278 test = butler.get(
279 "old_scarlet_model_data",
280 dataId={},
281 parameters={"blend_id": first_blend_id},
282 )
283 self.assertEqual(len(test.blends), 1)
284 self.assertIn(first_blend_id, test.blends)
286 def test_older_legacy_model(self):
287 repo = self._setup_butler()
288 oldStorageClass = StorageClass(
289 "ScarletModelData",
290 pytype=lsst.scarlet.lite.io.ScarletModelData,
291 )
292 oldDatasetType = DatasetType(
293 "old_scarlet_model_data",
294 dimensions=(),
295 storageClass=oldStorageClass,
296 universe=repo.dimensions,
297 )
298 ref = DatasetRef(
299 oldDatasetType,
300 run="test_ingestion",
301 dataId={},
302 )
303 dataset = FileDataset(
304 path=os.path.join(TESTDIR, "data", "v29_models.json"),
305 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
306 refs=[ref],
307 )
309 # Ingest the legacy model into the butler
310 butler = makeTestCollection(repo, uniqueId="ingestion")
311 repo.registry.registerDatasetType(oldDatasetType)
312 butler.ingest(dataset)
314 # Load the base repo config from the repository
315 base_config = Config(os.path.join(self.repo_dir, "butler.yaml"))
317 # Load the storage class override config
318 override_path = os.path.join(
319 os.path.dirname(lsst.daf.butler.__file__),
320 "configs",
321 "storageClasses.yaml"
322 )
323 override_config = Config(override_path)
325 # Merge the configs (update base with override)
326 base_config.update(override_config)
328 # Create Butler with the merged config
329 # The config now contains both the repo info and
330 # the storage class overrides
331 newButler = Butler.from_config(base_config, collections=butler.collections)
333 model = newButler.get("old_scarlet_model_data", dataId={}, storageClass="LsstScarletModelData")
334 self.assertEqual(len(model.blends), 2)
335 self.assertEqual(len(model.isolated), 0)
337 def test_read_legacy_zip_without_metadata(self):
338 """``read_scarlet_model`` reads a legacy-format zip that has no
339 ``metadata`` entry.
341 Legacy archives store the model PSF as top-level ``psf`` /
342 ``psfShape`` entries instead of a ``metadata`` entry.
343 ``zipfile.ZipFile.open`` raises ``KeyError`` (not ``ValueError``)
344 for a missing entry, so the legacy fallback was unreachable and
345 such archives crashed on read. Regression test for finding C-3
346 of the ``audits/audit-2026-05-05.md`` audit; also pins the IO-17
347 fix that the legacy load now produces a ``metadata['model_psf']``
348 numpy array.
349 """
350 bundle = pipeline.deblend(
351 pipeline.deconvolve(
352 pipeline.detect(pipeline.build_image(SCENES["multi-blend"]))
353 )
354 )
355 jm = bundle.result.scarletModelData.as_dict()
357 # Repackage the model in the legacy layout: one entry per blend
358 # plus a top-level model PSF, and crucially no ``metadata`` entry.
359 buf = io.BytesIO()
360 with zipfile.ZipFile(buf, "w") as zf:
361 for blendId, blendData in jm["blends"].items():
362 zf.writestr(str(blendId), json.dumps(blendData))
363 model_psf = jm["metadata"]["model_psf"]
364 model_psf_shape = list(np.asarray(model_psf).shape)
365 zf.writestr("psf", json.dumps(model_psf))
366 zf.writestr("psfShape", json.dumps(model_psf_shape))
367 buf.seek(0)
369 model = mes.io.utils.read_scarlet_model(buf)
370 self.assertEqual(len(model.blends), len(jm["blends"]))
371 self.assertIsNotNone(model.metadata)
372 self.assertIn("model_psf", model.metadata)
373 self.assertIsInstance(model.metadata["model_psf"], np.ndarray)
374 self.assertEqual(
375 list(model.metadata["model_psf"].shape), model_psf_shape
376 )
378 def _test_blend(self, blendData1, blendData2, model_psf, psf, bands):
379 # Test that two ScarletBlendData objects are equal
380 # up to machine precision.
381 self.assertTupleEqual(blendData1.origin, blendData2.origin)
382 self.assertEqual(len(blendData1.sources), len(blendData2.sources))
384 # Test that the two blends are equal up to machine precision
385 # once converted into scarlet lite Blend objects.
386 blend1 = blendData1.minimal_data_to_blend(
387 model_psf,
388 psf,
389 bands,
390 dtype=np.float32,
391 )
392 blend2 = blendData2.minimal_data_to_blend(
393 model_psf,
394 psf,
395 bands,
396 dtype=np.float32,
397 )
398 np.testing.assert_almost_equal(blend1.get_model().data, blend2.get_model().data)
400 def _setup_butler(self):
401 # Initialize a Butler to test persistence
402 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
403 self.repo_dir = repo_dir.name
404 self.addCleanup(tempfile.TemporaryDirectory.cleanup, repo_dir)
405 config = Config()
406 config["datastore", "cls"] = "lsst.daf.butler.datastores.fileDatastore.FileDatastore"
407 repo = makeTestRepo(repo_dir.name, config=config)
408 storageClass = StorageClass(
409 "LsstScarletModelData",
410 pytype=mes.io.LsstScarletModelData,
411 parameters=('blend_id',),
412 delegate="lsst.meas.extensions.scarlet.io.ScarletModelDelegate",
413 )
414 datasetType = DatasetType(
415 "scarlet_model_data",
416 dimensions=(),
417 storageClass=storageClass,
418 universe=repo.dimensions,
419 )
420 repo.registry.registerDatasetType(datasetType)
421 return repo
424def setup_module(module):
425 lsst.utils.tests.init()
428class MemoryTester(lsst.utils.tests.MemoryTestCase):
429 pass
432if __name__ == "__main__": 432 ↛ 433line 432 didn't jump to line 433 because the condition on line 432 was never true
433 lsst.utils.tests.init()
434 unittest.main()