Coverage for tests/test_io_persistence.py: 99%
190 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-20 09:21 +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.bands
80 self.model_psf = self.modelData.model_psf[None, :, :]
81 self.psf = self.modelData.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.model_psf[None, :, :], self.model_psf
99 )
100 np.testing.assert_almost_equal(modelData2.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 """A pre-``metadata`` (v29) archive loads and promotes its
167 ``psf`` / ``psfShape`` into the typed ``model_psf`` attribute.
169 """
170 model, butler = self._load_legacy_model("v29_models.json", "v29")
171 self.assertEqual(len(model.blends), 2)
172 self.assertNotIn("psfShape", model.metadata or {})
173 self._assert_single_blend_load(butler, 3495976385350991873)
175 def test_v30_legacy_model(self):
176 """``LsstScarletModelData`` ingested from a v30-era fixture
177 round-trips intact.
178 """
179 model, butler = self._load_legacy_model("v30_models.json", "v30")
181 # The multi-blend scene that generated the fixture produces
182 # three parent blends and one isolated source.
183 self.assertEqual(len(model.blends), 3)
184 self.assertEqual(len(model.isolated), 1)
186 # The per-band psf and band list round-trip as typed attributes.
187 self.assertIsNotNone(model.psf)
188 self.assertIsNotNone(model.bands)
190 # The isolated source survives the full ``IsolatedSourceData``
191 # round-trip: shape and integer peak (post-IO-1), and a
192 # bit-exact span mask. Pinning the span sum guards the
193 # ``span_array`` serialization path against silent regressions
194 # under future schema bumps.
195 iso = next(iter(model.isolated.values()))
196 self.assertEqual(iso.span_array.shape, (13, 13))
197 self.assertEqual(iso.origin, (6, 14))
198 self.assertEqual(iso.peak, (12, 20))
199 self.assertEqual(float(iso.span_array.sum()), 119.0)
201 # Single-blend parameter load also works on v30 archives.
202 self._assert_single_blend_load(butler, sorted(model.blends.keys())[0])
204 def test_v31a_legacy_model(self):
205 """A pre-DM-55109 (schema 1.0.1) archive promotes to the typed model.
207 """
208 model, butler = self._load_legacy_model("v31a_models.json", "v31a")
210 # The migration chain promoted the model to the current schema.
211 self.assertEqual(model.version, "1.0.2")
212 self.assertEqual(len(model.blends), 3)
213 self.assertEqual(len(model.isolated), 1)
215 # Model-level fields are now typed attributes.
216 self.assertEqual(tuple(model.bands), ("g", "r", "i"))
217 self.assertEqual(model.psf.shape, (3, 41, 41))
219 # Every parent became a typed blend; legacy_spans is False since the
220 # archive carried real footprint spans.
221 for blend in model.blends.values():
222 self.assertIsInstance(blend, mes.io.LsstHierarchicalBlendData)
223 self.assertFalse(blend.legacy_spans)
225 # Pin the first parent's promoted spans so the conversion stays
226 # bit-exact.
227 first_blend_id = sorted(model.blends.keys())[0]
228 first = model.blends[first_blend_id]
229 self.assertEqual(first.span_array.shape, (29, 41))
230 self.assertEqual(first.origin, (10, 50))
231 self.assertEqual(int(first.span_array.sum()), 797)
233 # The isolated source round-trips unchanged through the migration.
234 iso = next(iter(model.isolated.values()))
235 self.assertEqual(iso.span_array.shape, (13, 13))
236 self.assertEqual(iso.origin, (6, 14))
237 self.assertEqual(iso.peak, (12, 20))
238 self.assertEqual(float(iso.span_array.sum()), 119.0)
240 # Single-blend parameter load also works on v31a archives.
241 self._assert_single_blend_load(butler, first_blend_id)
243 def test_older_legacy_model(self):
244 repo = self._setup_butler()
245 oldStorageClass = StorageClass(
246 "ScarletModelData",
247 pytype=lsst.scarlet.lite.io.ScarletModelData,
248 )
249 oldDatasetType = DatasetType(
250 "old_scarlet_model_data",
251 dimensions=(),
252 storageClass=oldStorageClass,
253 universe=repo.dimensions,
254 )
255 ref = DatasetRef(
256 oldDatasetType,
257 run="test_ingestion",
258 dataId={},
259 )
260 dataset = FileDataset(
261 path=os.path.join(TESTDIR, "data", "v29_models.json"),
262 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
263 refs=[ref],
264 )
266 # Ingest the legacy model into the butler
267 butler = makeTestCollection(repo, uniqueId="ingestion")
268 repo.registry.registerDatasetType(oldDatasetType)
269 butler.ingest(dataset)
271 # Load the base repo config from the repository
272 base_config = Config(os.path.join(self.repo_dir, "butler.yaml"))
274 # Load the storage class override config
275 override_path = os.path.join(
276 os.path.dirname(lsst.daf.butler.__file__),
277 "configs",
278 "storageClasses.yaml"
279 )
280 override_config = Config(override_path)
282 # Merge the configs (update base with override)
283 base_config.update(override_config)
285 # Create Butler with the merged config
286 # The config now contains both the repo info and
287 # the storage class overrides
288 newButler = Butler.from_config(base_config, collections=butler.collections)
290 model = newButler.get("old_scarlet_model_data", dataId={}, storageClass="LsstScarletModelData")
291 self.assertEqual(len(model.blends), 2)
292 self.assertEqual(len(model.isolated), 0)
294 def test_lsst_scarlet_model_conversion(self):
295 """Test converting an LsstScarletModelData to a ScarletModelData
296 (lossily) via the Butler.
297 """
298 repo = self._setup_butler()
299 newDatasetType = DatasetType(
300 "new_scarlet_model_data",
301 dimensions=(),
302 storageClass="LsstScarletModelData",
303 universe=repo.dimensions,
304 )
305 ref = DatasetRef(
306 newDatasetType,
307 run="test_ingestion",
308 dataId={},
309 )
310 dataset = FileDataset(
311 path=os.path.join(TESTDIR, "data", "v29_models.json"),
312 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
313 refs=[ref],
314 )
315 butler = makeTestCollection(repo, uniqueId="ingestion")
316 repo.registry.registerDatasetType(newDatasetType)
317 butler.ingest(dataset)
318 model = butler.get("new_scarlet_model_data", dataId={}, storageClass="ScarletModelData")
319 self.assertIsInstance(model, lsst.scarlet.lite.io.ScarletModelData)
320 self.assertEqual(len(model.blends), 2)
322 def test_read_legacy_zip_without_metadata(self):
323 """``read_scarlet_model`` reads a legacy-format zip that has no
324 ``metadata`` entry.
326 Legacy archives store the model PSF as top-level ``psf`` /
327 ``psfShape`` entries instead of a ``metadata`` entry.
328 ``zipfile.ZipFile.open`` raises ``KeyError`` (not ``ValueError``)
329 for a missing entry, so the legacy fallback was unreachable and
330 such archives crashed on read. Regression test for finding C-3
331 of the ``audits/audit-2026-05-05.md`` audit; also pins the IO-17
332 fix that the legacy load now produces a ``metadata['model_psf']``
333 numpy array.
334 """
335 bundle = pipeline.deblend(
336 pipeline.deconvolve(
337 pipeline.detect(pipeline.build_image(SCENES["multi-blend"]))
338 )
339 )
340 jm = bundle.result.scarletModelData.as_dict()
342 # Repackage the model in the legacy layout: one entry per blend
343 # plus a top-level model PSF, and crucially no ``metadata`` entry.
344 buf = io.BytesIO()
345 with zipfile.ZipFile(buf, "w") as zf:
346 for blendId, blendData in jm["blends"].items():
347 zf.writestr(str(blendId), json.dumps(blendData))
348 model_psf = jm["metadata"]["model_psf"]
349 model_psf_shape = list(np.asarray(model_psf).shape)
350 zf.writestr("psf", json.dumps(model_psf))
351 zf.writestr("psfShape", json.dumps(model_psf_shape))
352 buf.seek(0)
354 model = mes.io.utils.read_scarlet_model(buf)
355 self.assertEqual(len(model.blends), len(jm["blends"]))
356 self.assertIsNotNone(model.model_psf)
357 self.assertIsInstance(model.model_psf, np.ndarray)
358 self.assertEqual(
359 list(model.model_psf.shape), model_psf_shape
360 )
362 def _test_blend(self, blendData1, blendData2, model_psf, psf, bands):
363 # Test that two ScarletBlendData objects are equal
364 # up to machine precision.
365 self.assertTupleEqual(blendData1.origin, blendData2.origin)
366 self.assertEqual(len(blendData1.sources), len(blendData2.sources))
368 # Test that the two blends are equal up to machine precision
369 # once converted into scarlet lite Blend objects.
370 blend1 = blendData1.minimal_data_to_blend(
371 model_psf,
372 psf,
373 bands,
374 dtype=np.float32,
375 )
376 blend2 = blendData2.minimal_data_to_blend(
377 model_psf,
378 psf,
379 bands,
380 dtype=np.float32,
381 )
382 np.testing.assert_almost_equal(blend1.get_model().data, blend2.get_model().data)
384 def _load_legacy_model(self, filename, unique):
385 """Ingest a legacy JSON model test context and return
386 ``(model, butler)``.
388 Parameters
389 ----------
390 filename : str
391 Fixture name under ``tests/data``.
392 unique : str
393 Short tag making the ingestion run/collection names unique.
394 """
395 repo = self._setup_butler()
396 storageClass = StorageClass(
397 "LsstScarletModelData",
398 pytype=mes.io.LsstScarletModelData,
399 )
400 datasetType = DatasetType(
401 "old_scarlet_model_data",
402 dimensions=(),
403 storageClass=storageClass,
404 universe=repo.dimensions,
405 )
406 ref = DatasetRef(datasetType, run=f"test_ingestion_{unique}", dataId={})
407 dataset = FileDataset(
408 path=os.path.join(TESTDIR, "data", filename),
409 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
410 refs=[ref],
411 )
413 butler = makeTestCollection(repo, uniqueId=f"ingestion_{unique}")
414 repo.registry.registerDatasetType(datasetType)
415 butler.ingest(dataset)
417 model = butler.get("old_scarlet_model_data", dataId={})
418 self.assertIsInstance(model.model_psf, np.ndarray)
419 self.assertEqual(model.model_psf.shape, (15, 15))
420 return model, butler
422 def _assert_single_blend_load(self, butler, blend_id):
423 """A ``blend_id`` parameter load returns exactly that one blend."""
424 test = butler.get(
425 "old_scarlet_model_data",
426 dataId={},
427 parameters={"blend_id": blend_id},
428 )
429 self.assertEqual(len(test.blends), 1)
430 self.assertIn(blend_id, test.blends)
432 def _setup_butler(self):
433 # Initialize a Butler to test persistence
434 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
435 self.repo_dir = repo_dir.name
436 self.addCleanup(tempfile.TemporaryDirectory.cleanup, repo_dir)
437 config = Config()
438 config["datastore", "cls"] = "lsst.daf.butler.datastores.fileDatastore.FileDatastore"
439 repo = makeTestRepo(repo_dir.name, config=config)
440 storageClass = StorageClass(
441 "LsstScarletModelData",
442 pytype=mes.io.LsstScarletModelData,
443 parameters=('blend_id',),
444 delegate="lsst.meas.extensions.scarlet.io.ScarletModelDelegate",
445 )
446 datasetType = DatasetType(
447 "scarlet_model_data",
448 dimensions=(),
449 storageClass=storageClass,
450 universe=repo.dimensions,
451 )
452 repo.registry.registerDatasetType(datasetType)
453 return repo
456def setup_module(module):
457 lsst.utils.tests.init()
460class MemoryTester(lsst.utils.tests.MemoryTestCase):
461 pass
464if __name__ == "__main__": 464 ↛ 465line 464 didn't jump to line 465 because the condition on line 464 was never true
465 lsst.utils.tests.init()
466 unittest.main()