Coverage for tests/test_io_persistence.py: 99%
194 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-09-14 02:43 -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/>.
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_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 oldDatasetType = DatasetType(
304 "older_scarlet_model_data",
305 dimensions=(),
306 storageClass="ScarletModelData",
307 universe=butler.dimensions,
308 )
309 butler.registry.registerDatasetType(oldDatasetType)
310 # Write the new model with the old storage class, which should be
311 # lossless.
312 butler.put(model1, "older_scarlet_model_data", dataId={})
313 # Read it back as the new (full) type, which also should be lossless.
314 model2 = butler.get(
315 "older_scarlet_model_data", dataId={}, storageClass="LsstScarletModelData"
316 )
317 self.assertIsInstance(model2, mes.io.LsstScarletModelData)
318 self.assertEqual(len(model2.blends), len(model1.blends))
319 self.assertEqual(set(model2.isolated), set(model1.isolated))
320 for sourceId in model1.isolated:
321 iso1 = model1.isolated[sourceId]
322 iso2 = model2.isolated[sourceId]
323 self.assertTupleEqual(iso1.origin, iso2.origin)
324 np.testing.assert_array_equal(iso1.span_array, iso2.span_array)
326 def test_read_legacy_zip_without_metadata(self):
327 """``read_scarlet_model`` reads a legacy-format zip that has no
328 ``metadata`` entry.
330 Legacy archives store the model PSF as top-level ``psf`` /
331 ``psfShape`` entries instead of a ``metadata`` entry.
332 ``zipfile.ZipFile.open`` raises ``KeyError`` (not ``ValueError``)
333 for a missing entry, so the legacy fallback was unreachable and
334 such archives crashed on read. Regression test for finding C-3
335 of the ``audits/audit-2026-05-05.md`` audit; also pins the IO-17
336 fix that the legacy load now produces a ``metadata['model_psf']``
337 numpy array.
338 """
339 bundle = pipeline.deblend(
340 pipeline.deconvolve(
341 pipeline.detect(pipeline.build_image(SCENES["multi-blend"]))
342 )
343 )
344 jm = bundle.result.scarletModelData.as_dict()
346 # Repackage the model in the legacy layout: one entry per blend
347 # plus a top-level model PSF, and crucially no ``metadata`` entry.
348 buf = io.BytesIO()
349 with zipfile.ZipFile(buf, "w") as zf:
350 for blendId, blendData in jm["blends"].items():
351 zf.writestr(str(blendId), json.dumps(blendData))
352 model_psf = jm["metadata"]["model_psf"]
353 model_psf_shape = list(np.asarray(model_psf).shape)
354 zf.writestr("psf", json.dumps(model_psf))
355 zf.writestr("psfShape", json.dumps(model_psf_shape))
356 buf.seek(0)
358 model = mes.io.utils.read_scarlet_model(buf)
359 self.assertEqual(len(model.blends), len(jm["blends"]))
360 self.assertIsNotNone(model.model_psf)
361 self.assertIsInstance(model.model_psf, np.ndarray)
362 self.assertEqual(
363 list(model.model_psf.shape), model_psf_shape
364 )
366 def _test_blend(self, blendData1, blendData2, model_psf, psf, bands):
367 # Test that two ScarletBlendData objects are equal
368 # up to machine precision.
369 self.assertTupleEqual(blendData1.origin, blendData2.origin)
370 self.assertEqual(len(blendData1.sources), len(blendData2.sources))
372 # Test that the two blends are equal up to machine precision
373 # once converted into scarlet lite Blend objects.
374 blend1 = blendData1.minimal_data_to_blend(
375 model_psf,
376 psf,
377 bands,
378 dtype=np.float32,
379 )
380 blend2 = blendData2.minimal_data_to_blend(
381 model_psf,
382 psf,
383 bands,
384 dtype=np.float32,
385 )
386 np.testing.assert_almost_equal(blend1.get_model().data, blend2.get_model().data)
388 def _load_legacy_model(self, filename, unique):
389 """Ingest a legacy JSON model test context and return
390 ``(model, butler)``.
392 Parameters
393 ----------
394 filename : str
395 Fixture name under ``tests/data``.
396 unique : str
397 Short tag making the ingestion run/collection names unique.
398 """
399 repo = self._setup_butler()
400 storageClass = StorageClass(
401 "LsstScarletModelData",
402 pytype=mes.io.LsstScarletModelData,
403 )
404 datasetType = DatasetType(
405 "old_scarlet_model_data",
406 dimensions=(),
407 storageClass=storageClass,
408 universe=repo.dimensions,
409 )
410 ref = DatasetRef(datasetType, run=f"test_ingestion_{unique}", dataId={})
411 dataset = FileDataset(
412 path=os.path.join(TESTDIR, "data", filename),
413 formatter="lsst.daf.butler.formatters.json.JsonFormatter",
414 refs=[ref],
415 )
417 butler = makeTestCollection(repo, uniqueId=f"ingestion_{unique}")
418 repo.registry.registerDatasetType(datasetType)
419 butler.ingest(dataset)
421 model = butler.get("old_scarlet_model_data", dataId={})
422 self.assertIsInstance(model.model_psf, np.ndarray)
423 self.assertEqual(model.model_psf.shape, (15, 15))
424 return model, butler
426 def _assert_single_blend_load(self, butler, blend_id):
427 """A ``blend_id`` parameter load returns exactly that one blend."""
428 test = butler.get(
429 "old_scarlet_model_data",
430 dataId={},
431 parameters={"blend_id": blend_id},
432 )
433 self.assertEqual(len(test.blends), 1)
434 self.assertIn(blend_id, test.blends)
436 def _setup_butler(self):
437 # Initialize a Butler to test persistence
438 repo_dir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
439 self.repo_dir = repo_dir.name
440 self.addCleanup(tempfile.TemporaryDirectory.cleanup, repo_dir)
441 config = Config()
442 config["datastore", "cls"] = "lsst.daf.butler.datastores.fileDatastore.FileDatastore"
443 repo = makeTestRepo(repo_dir.name, config=config)
444 storageClass = StorageClass(
445 "LsstScarletModelData",
446 pytype=mes.io.LsstScarletModelData,
447 parameters=('blend_id',),
448 delegate="lsst.meas.extensions.scarlet.io.ScarletModelDelegate",
449 )
450 datasetType = DatasetType(
451 "scarlet_model_data",
452 dimensions=(),
453 storageClass=storageClass,
454 universe=repo.dimensions,
455 )
456 repo.registry.registerDatasetType(datasetType)
457 return repo
460def setup_module(module):
461 lsst.utils.tests.init()
464class MemoryTester(lsst.utils.tests.MemoryTestCase):
465 pass
468if __name__ == "__main__": 468 ↛ 469line 468 didn't jump to line 469 because the condition on line 468 was never true
469 lsst.utils.tests.init()
470 unittest.main()