Coverage for tests/test_io_persistence.py: 99%
210 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-24 09:12 +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.defaults)
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_write_conversion(self):
295 """Test writing an LsstScarletModelData to a dataset type with
296 the old ScarletModelData storage class via the Butler.
297 """
298 model1, butler = self._load_legacy_model(
299 os.path.join(TESTDIR, "data", "v31a_models.json"),
300 "write_conversion",
301 )
302 self.assertIsInstance(model1, mes.io.LsstScarletModelData)
303 old_types = {}
304 for name in ["put_by_name", "put_by_type", "put_by_old_ref", "put_by_new_ref"]:
305 with self.subTest(name):
306 old_type = DatasetType(
307 name,
308 dimensions=(),
309 storageClass="ScarletModelData",
310 universe=butler.dimensions,
311 )
312 butler.registry.registerDatasetType(old_type)
313 old_types[name] = old_type
315 # Conversion in this direction needs to look like it can work
316 # to give the formatter a chance to handle an
317 # LsstScarletModelInstance.
318 # for-loop variable leakage is intentional here; we just need to test
319 # one of these.
320 new_type = old_type.overrideStorageClass("LsstScarletModelData")
321 self.assertTrue(old_type.is_compatible_with(new_type))
322 self.assertTrue(old_type.storageClass.can_convert(new_type.storageClass))
323 # The other direction is backed by the real read converter.
324 self.assertTrue(new_type.is_compatible_with(old_type))
326 # Write the new model in three different ways with the old storage
327 # class, which should be lossless.
328 butler.put(model1, "put_by_name", dataId={})
329 butler.put(model1, old_types["put_by_type"], dataId={})
330 butler.put(model1, DatasetRef(old_types["put_by_old_ref"], run=butler.run, dataId={}))
331 butler.put(
332 model1,
333 DatasetRef(
334 old_types["put_by_new_ref"], run=butler.run, dataId={}
335 ).overrideStorageClass("LsstScarletModelData"),
336 )
338 # Read each back as the new (full) type, which also should be
339 # lossless.
340 for name in old_types:
341 model2 = butler.get(
342 name, dataId={}, storageClass="LsstScarletModelData"
343 )
344 self.assertIsInstance(model2, mes.io.LsstScarletModelData)
345 self.assertEqual(len(model2.blends), len(model1.blends))
346 self.assertEqual(set(model2.isolated), set(model1.isolated))
347 for sourceId in model1.isolated:
348 iso1 = model1.isolated[sourceId]
349 iso2 = model2.isolated[sourceId]
350 self.assertTupleEqual(iso1.origin, iso2.origin)
351 np.testing.assert_array_equal(iso1.span_array, iso2.span_array)
352 # Reading without a storage class override cannot produce the old
353 # (never-written) type; the error should suggest how to read it.
354 with self.assertRaises(TypeError) as raisesCtx:
355 butler.get(name, dataId={})
356 self.assertIn("storageClass", str(raisesCtx.exception))
357 self.assertIn("LsstScarletModelData", str(raisesCtx.exception))
359 def test_read_legacy_zip_without_metadata(self):
360 """``read_scarlet_model`` reads a legacy-format zip that has no
361 ``metadata`` entry.
363 Legacy archives store the model PSF as top-level ``psf`` /
364 ``psfShape`` entries instead of a ``metadata`` entry.
365 ``zipfile.ZipFile.open`` raises ``KeyError`` (not ``ValueError``)
366 for a missing entry, so the legacy fallback was unreachable and
367 such archives crashed on read. Regression test for finding C-3
368 of the ``audits/audit-2026-05-05.md`` audit; also pins the IO-17
369 fix that the legacy load now produces a ``metadata['model_psf']``
370 numpy array.
371 """
372 bundle = pipeline.deblend(
373 pipeline.deconvolve(
374 pipeline.detect(pipeline.build_image(SCENES["multi-blend"]))
375 )
376 )
377 jm = bundle.result.scarletModelData.as_dict()
379 # Repackage the model in the legacy layout: one entry per blend
380 # plus a top-level model PSF, and crucially no ``metadata`` entry.
381 buf = io.BytesIO()
382 with zipfile.ZipFile(buf, "w") as zf:
383 for blendId, blendData in jm["blends"].items():
384 zf.writestr(str(blendId), json.dumps(blendData))
385 model_psf = jm["metadata"]["model_psf"]
386 model_psf_shape = list(np.asarray(model_psf).shape)
387 zf.writestr("psf", json.dumps(model_psf))
388 zf.writestr("psfShape", json.dumps(model_psf_shape))
389 buf.seek(0)
391 model = mes.io.utils.read_scarlet_model(buf)
392 self.assertEqual(len(model.blends), len(jm["blends"]))
393 self.assertIsNotNone(model.model_psf)
394 self.assertIsInstance(model.model_psf, np.ndarray)
395 self.assertEqual(
396 list(model.model_psf.shape), model_psf_shape
397 )
399 def _test_blend(self, blendData1, blendData2, model_psf, psf, bands):
400 # Test that two ScarletBlendData objects are equal
401 # up to machine precision.
402 self.assertTupleEqual(blendData1.origin, blendData2.origin)
403 self.assertEqual(len(blendData1.sources), len(blendData2.sources))
405 # Test that the two blends are equal up to machine precision
406 # once converted into scarlet lite Blend objects.
407 blend1 = blendData1.minimal_data_to_blend(
408 model_psf,
409 psf,
410 bands,
411 dtype=np.float32,
412 )
413 blend2 = blendData2.minimal_data_to_blend(
414 model_psf,
415 psf,
416 bands,
417 dtype=np.float32,
418 )
419 np.testing.assert_almost_equal(blend1.get_model().data, blend2.get_model().data)
421 def _load_legacy_model(self, filename, unique):
422 """Ingest a legacy JSON model test context and return
423 ``(model, butler)``.
425 Parameters
426 ----------
427 filename : str
428 Fixture name under ``tests/data``.
429 unique : str
430 Short tag making the ingestion run/collection names unique.
431 """
432 repo = self._setup_butler()
433 storageClass = StorageClass(
434 "LsstScarletModelData",
435 pytype=mes.io.LsstScarletModelData,
436 )
437 datasetType = DatasetType(
438 "old_scarlet_model_data",
439 dimensions=(),
440 storageClass=storageClass,
441 universe=repo.dimensions,
442 )
443 ref = DatasetRef(datasetType, run=f"test_ingestion_{unique}", dataId={})
444 dataset = FileDataset(
445 path=os.path.join(TESTDIR, "data", filename),
446 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
447 refs=[ref],
448 )
450 butler = makeTestCollection(repo, uniqueId=f"ingestion_{unique}")
451 repo.registry.registerDatasetType(datasetType)
452 butler.ingest(dataset)
454 model = butler.get("old_scarlet_model_data", dataId={})
455 self.assertIsInstance(model.model_psf, np.ndarray)
456 self.assertEqual(model.model_psf.shape, (15, 15))
457 return model, butler
459 def _assert_single_blend_load(self, butler, blend_id):
460 """A ``blend_id`` parameter load returns exactly that one blend."""
461 test = butler.get(
462 "old_scarlet_model_data",
463 dataId={},
464 parameters={"blend_id": blend_id},
465 )
466 self.assertEqual(len(test.blends), 1)
467 self.assertIn(blend_id, test.blends)
469 def _setup_butler(self):
470 # Initialize a Butler to test persistence
471 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
472 self.repo_dir = repo_dir.name
473 self.addCleanup(tempfile.TemporaryDirectory.cleanup, repo_dir)
474 config = Config()
475 config["datastore", "cls"] = "lsst.daf.butler.datastores.fileDatastore.FileDatastore"
476 repo = makeTestRepo(repo_dir.name, config=config)
477 storageClass = StorageClass(
478 "LsstScarletModelData",
479 pytype=mes.io.LsstScarletModelData,
480 parameters=('blend_id',),
481 delegate="lsst.meas.extensions.scarlet.io.ScarletModelDelegate",
482 )
483 datasetType = DatasetType(
484 "scarlet_model_data",
485 dimensions=(),
486 storageClass=storageClass,
487 universe=repo.dimensions,
488 )
489 repo.registry.registerDatasetType(datasetType)
490 return repo
493def setup_module(module):
494 lsst.utils.tests.init()
497class MemoryTester(lsst.utils.tests.MemoryTestCase):
498 pass
501if __name__ == "__main__": 501 ↛ 502line 501 didn't jump to line 502 because the condition on line 501 was never true
502 lsst.utils.tests.init()
503 unittest.main()