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