Coverage for python/lsst/images/tests/extract_legacy_test_data.py: 21%
98 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-22 01:54 -0700
1# This file is part of lsst-images.
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# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = ()
16import os
17from typing import TYPE_CHECKING
19import click
20import numpy as np
22if TYPE_CHECKING:
23 from lsst.daf.butler import Butler, DatasetRef
26from ._data_ids import DP2_COADD_DATA_ID, DP2_COADD_MISSING_CELL, DP2_VISIT_DETECTOR_DATA_ID
29def extract_exposure(
30 butler: Butler,
31 output_path: str,
32 dataset_ref: DatasetRef,
33 shuffle: bool,
34) -> None:
35 """Load a subimage of a processed visit image from a butler repository
36 and save it to testdata_images.
37 """
38 from lsst.afw.fits import (
39 CompressionAlgorithm,
40 CompressionOptions,
41 DitherAlgorithm,
42 QuantizationOptions,
43 ScalingAlgorithm,
44 )
45 from lsst.geom import Box2I, Extent2I, Point2I
47 exposure = butler.get(dataset_ref, parameters={"bbox": Box2I(Point2I(5, 4), Extent2I(256, 250))})
48 if shuffle:
49 indices = np.arange(exposure.image.array.size, dtype=int)
50 rng = np.random.default_rng()
51 rng.shuffle(indices)
52 exposure.image.array[:, :] = exposure.image.array.flat[indices].reshape(250, 256)
53 exposure.mask.array[:, :] = exposure.mask.array.flat[indices].reshape(250, 256)
54 exposure.variance.array[:, :] = exposure.variance.array.flat[indices].reshape(250, 256)
55 float_compression = CompressionOptions(
56 algorithm=CompressionAlgorithm.RICE_1,
57 tile_height=50,
58 tile_width=64,
59 quantization=QuantizationOptions(
60 dither=DitherAlgorithm.SUBTRACTIVE_DITHER_2,
61 scaling=ScalingAlgorithm.STDEV_MASKED,
62 level=16,
63 seed=747,
64 ),
65 )
66 mask_compression = CompressionOptions(
67 algorithm=CompressionAlgorithm.GZIP_2,
68 tile_height=50,
69 tile_width=64,
70 quantization=None,
71 )
72 os.makedirs(os.path.dirname(output_path), exist_ok=True)
73 exposure.writeFits(
74 output_path,
75 imageOptions=float_compression,
76 maskOptions=mask_compression,
77 varianceOptions=float_compression,
78 )
81def extract_visit_image_background(
82 butler: Butler,
83 output_path: str,
84 dataset_ref: DatasetRef,
85) -> None:
86 """Load the background model of a processed visit image from a butler
87 repository and save it to testdata_images.
88 """
89 visit_image_background = butler.get(dataset_ref)
90 visit_image_background.writeFits(output_path)
93def extract_visit_summary(
94 butler: Butler,
95 output_path: str,
96 dataset_ref: DatasetRef,
97) -> None:
98 """Load a visit_summary dataset and save it to testdata_images."""
99 visit_summary = butler.get(dataset_ref)
100 visit_summary.writeFits(output_path)
103def extract_cell_coadd(
104 butler: Butler,
105 output_path: str,
106 dataset_ref: DatasetRef,
107 shuffle: bool,
108) -> None:
109 """Load a subimage of a cell coadd from a butler repository and save it
110 to testdata_images.
111 """
112 from lsst.cell_coadds import MultipleCellCoadd
114 full_cell_coadd = butler.get(dataset_ref)
115 cell_coadd = MultipleCellCoadd(
116 [
117 full_cell_coadd.cells[x, y]
118 for y in range(7, 11)
119 for x in range(5, 8)
120 if {"i": y, "j": x} != DP2_COADD_MISSING_CELL
121 ],
122 grid=full_cell_coadd.grid,
123 outer_cell_size=full_cell_coadd.outer_cell_size,
124 psf_image_size=full_cell_coadd.psf_image_size,
125 common=full_cell_coadd.common,
126 )
127 if shuffle:
128 rng = np.random.default_rng()
129 for cell in cell_coadd.cells.values():
130 indices = np.arange(cell.outer.image.array.size, dtype=int)
131 rng.shuffle(indices)
132 cell.outer.image.array[:, :] = cell.outer.image.array.flat[indices].reshape(150, 150)
133 cell.outer.mask.array[:, :] = cell.outer.mask.array.flat[indices].reshape(150, 150)
134 cell.outer.variance.array[:, :] = cell.outer.variance.array.flat[indices].reshape(150, 150)
135 for n in cell.outer.noise_realizations:
136 n.array[:, :] = n.array.flat[indices].reshape(150, 150)
137 if cell.outer.mask_fractions is not None:
138 cell.outer.mask_fractions.array[:, :] = cell.outer.mask_fractions.array.flat[indices].reshape(
139 150, 150
140 )
141 os.makedirs(os.path.dirname(output_path), exist_ok=True)
142 cell_coadd.writeFits(output_path)
145def extract_camera(butler: Butler, output_path: str, dataset_ref: DatasetRef) -> None:
146 """Read camera geometry from a butler repository and save it to
147 testdata_images.
148 """
149 camera = butler.get(dataset_ref)
150 os.makedirs(os.path.dirname(output_path), exist_ok=True)
151 camera.writeFits(output_path)
154def extract_skymap(butler: Butler, output_path: str, dataset_ref: DatasetRef) -> None:
155 """Read a skymap definition from a butler repository and save it to
156 testdata_images.
157 """
158 os.makedirs(os.path.dirname(output_path), exist_ok=True)
159 (path,) = butler.retrieveArtifacts(
160 [dataset_ref],
161 destination=os.path.dirname(output_path),
162 transfer="copy",
163 preserve_path=False,
164 overwrite=True,
165 )
166 if path.ospath != output_path:
167 os.rename(path.ospath, output_path)
170def find_dataset_or_raise(
171 butler: Butler, dataset_type: str, *, collections: str | None = None, **kwargs
172) -> DatasetRef:
173 """Call `lsst.daf.butler.Butler.find_dataset` with the given arguments and
174 raise `LookupError` if it returns `None`.
175 """
176 ref = butler.find_dataset(dataset_type, collections=collections, **kwargs)
177 if ref is None:
178 raise LookupError(f"Could not find dataset {dataset_type} with data ID {kwargs}.")
179 return ref
182@click.group("extract_test_data")
183def extract_test_data() -> None:
184 """Extract test fixtures from a Rubin data repository."""
187@extract_test_data.command("dp2")
188@click.option("-b", "--butler-repo", help="Path to the butler repository.")
189@click.option("-d", "--testdata-dir", help="Path to the testdata_images directory.")
190@click.option(
191 "-c",
192 "--collection",
193 default="LSSTCam/runs/DRP/DP2",
194 help="Collection to use for most data products.",
195)
196@click.option(
197 "--visit-images/--no-visit-images",
198 default=True,
199 help="Whether to extract [preliminary_]visit_image datasets.",
200)
201@click.option(
202 "--difference-images/--no-difference-images",
203 default=True,
204 help="Whether to extract difference_image datasets.",
205)
206@click.option(
207 "--coadds/--no-coadds",
208 default=True,
209 help="Whether to extract coadd datasets.",
210)
211@click.option(
212 "--camera/--no-camera",
213 default=True,
214 help="Whether to extract the camera.",
215)
216@click.option(
217 "--skymap/--no-skymap",
218 default=True,
219 help="Whether to extract the skymap.",
220)
221def extract_dp2(
222 butler_repo: str | None,
223 testdata_dir: str | None,
224 collection: str,
225 *,
226 visit_images: bool,
227 difference_images: bool,
228 coadds: bool,
229 camera: bool,
230 skymap: bool,
231) -> None:
232 """Extract test data from a butler repository."""
233 try:
234 # lsst.afw.image is imported only to confirm a full Rubin pipelines
235 # environment is present: daf.butler and lsst.utils are available
236 # from the pip-installed package and so cannot gate on their own.
237 import lsst.afw.image # noqa: F401
238 from lsst.daf.butler import Butler
239 from lsst.utils import getPackageDir
240 except ImportError as err:
241 err.add_note(
242 "Updating the test data requires a full Rubin development enviroment with at least "
243 "'afw', 'obs_base', 'meas_extensions_psfex', 'meas_extensions_piff' and 'cell_coadds' "
244 "importable. This is not necessary for just running the tests."
245 )
246 raise
247 if butler_repo is None:
248 butler_repo = "dp2_prep"
249 if testdata_dir is None:
250 testdata_dir = getPackageDir("testdata_images")
251 butler = Butler.from_config(butler_repo, collections=[collection])
252 if visit_images:
253 extract_exposure(
254 butler,
255 os.path.join(testdata_dir, "dp2", "legacy", "visit_image.fits"),
256 find_dataset_or_raise(butler, "visit_image", **DP2_VISIT_DETECTOR_DATA_ID),
257 shuffle=True,
258 )
259 extract_visit_image_background(
260 butler,
261 os.path.join(testdata_dir, "dp2", "legacy", "visit_image_background.fits"),
262 find_dataset_or_raise(butler, "visit_image_background", **DP2_VISIT_DETECTOR_DATA_ID),
263 )
264 extract_visit_summary(
265 butler,
266 os.path.join(testdata_dir, "dp2", "legacy", "visit_summary.fits"),
267 find_dataset_or_raise(butler, "visit_summary", **DP2_VISIT_DETECTOR_DATA_ID),
268 )
269 extract_exposure(
270 butler,
271 os.path.join(testdata_dir, "dp2", "legacy", "preliminary_visit_image.fits"),
272 find_dataset_or_raise(butler, "preliminary_visit_image", **DP2_VISIT_DETECTOR_DATA_ID),
273 shuffle=True,
274 )
275 extract_visit_image_background(
276 butler,
277 os.path.join(testdata_dir, "dp2", "legacy", "preliminary_visit_image_background.fits"),
278 find_dataset_or_raise(butler, "preliminary_visit_image_background", **DP2_VISIT_DETECTOR_DATA_ID),
279 )
280 if difference_images:
281 extract_exposure(
282 butler,
283 os.path.join(testdata_dir, "dp2", "legacy", "difference_image.fits"),
284 find_dataset_or_raise(butler, "difference_image", **DP2_VISIT_DETECTOR_DATA_ID),
285 shuffle=True,
286 )
287 if coadds:
288 extract_cell_coadd(
289 butler,
290 os.path.join(
291 testdata_dir,
292 "dp2",
293 "legacy",
294 "deep_coadd_cell_predetection.fits",
295 ),
296 find_dataset_or_raise(butler, "deep_coadd_cell_predetection", **DP2_COADD_DATA_ID),
297 shuffle=True,
298 )
299 if skymap:
300 extract_skymap(
301 butler,
302 os.path.join(
303 testdata_dir,
304 "dp2",
305 "legacy",
306 "skyMap.pickle",
307 ),
308 find_dataset_or_raise(butler, "skyMap", skymap=DP2_COADD_DATA_ID["skymap"]),
309 )
310 if camera:
311 extract_camera(
312 butler,
313 os.path.join(
314 testdata_dir,
315 "dp2",
316 "legacy",
317 "camera.fits",
318 ),
319 find_dataset_or_raise(butler, "camera", instrument="LSSTCam"),
320 )
323if __name__ == "__main__":
324 extract_test_data()