Coverage for python/lsst/images/tests/extract_legacy_test_data.py: 20%
104 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-31 09:23 +0000
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, Any
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.
38 Parameters
39 ----------
40 butler
41 Butler repository to read the dataset from.
42 output_path
43 Path to the FITS file to write under testdata_images.
44 dataset_ref
45 Reference to the dataset to load from the butler.
46 shuffle
47 If `True`, randomly permute the image, mask, and variance pixels so
48 the test data does not reproduce real survey pixel values.
49 """
50 from lsst.afw.fits import (
51 CompressionAlgorithm,
52 CompressionOptions,
53 DitherAlgorithm,
54 QuantizationOptions,
55 ScalingAlgorithm,
56 )
57 from lsst.afw.image import LOCAL
58 from lsst.geom import Box2I, Extent2I, Point2I
60 exposure = butler.get(
61 dataset_ref,
62 # Interpret bbox as LOCAL to let us select from a patch-level coadd
63 # image without knowing that patch's offset (LOCAL==PARENT for
64 # detector-level images).
65 parameters={"bbox": Box2I(Point2I(5, 4), Extent2I(256, 250)), "origin": LOCAL},
66 )
67 if shuffle:
68 indices = np.arange(exposure.image.array.size, dtype=int)
69 rng = np.random.default_rng()
70 rng.shuffle(indices)
71 exposure.image.array[:, :] = exposure.image.array.flat[indices].reshape(250, 256)
72 exposure.mask.array[:, :] = exposure.mask.array.flat[indices].reshape(250, 256)
73 exposure.variance.array[:, :] = exposure.variance.array.flat[indices].reshape(250, 256)
74 float_compression = CompressionOptions(
75 algorithm=CompressionAlgorithm.RICE_1,
76 tile_height=50,
77 tile_width=64,
78 quantization=QuantizationOptions(
79 dither=DitherAlgorithm.SUBTRACTIVE_DITHER_2,
80 scaling=ScalingAlgorithm.STDEV_MASKED,
81 level=16,
82 seed=747,
83 ),
84 )
85 mask_compression = CompressionOptions(
86 algorithm=CompressionAlgorithm.GZIP_2,
87 tile_height=50,
88 tile_width=64,
89 quantization=None,
90 )
91 os.makedirs(os.path.dirname(output_path), exist_ok=True)
92 exposure.writeFits(
93 output_path,
94 imageOptions=float_compression,
95 maskOptions=mask_compression,
96 varianceOptions=float_compression,
97 )
100def extract_visit_image_background(
101 butler: Butler,
102 output_path: str,
103 dataset_ref: DatasetRef,
104) -> None:
105 """Load the background model of a processed visit image from a butler
106 repository and save it to testdata_images.
108 Parameters
109 ----------
110 butler
111 Butler repository to read the dataset from.
112 output_path
113 Path to the FITS file to write under testdata_images.
114 dataset_ref
115 Reference to the dataset to load from the butler.
116 """
117 visit_image_background = butler.get(dataset_ref)
118 visit_image_background.writeFits(output_path)
121def extract_generic(
122 butler: Butler,
123 output_path: str,
124 dataset_ref: DatasetRef,
125) -> None:
126 """Load a dataset and save it to testdata_images, assuming it has a
127 writeFits method.
129 Parameters
130 ----------
131 butler
132 Butler repository to read the dataset from.
133 output_path
134 Path to the FITS file to write under testdata_images.
135 dataset_ref
136 Reference to the dataset to load from the butler.
137 """
138 visit_summary = butler.get(dataset_ref)
139 visit_summary.writeFits(output_path)
142def extract_cell_coadd(
143 butler: Butler,
144 output_path: str,
145 dataset_ref: DatasetRef,
146 shuffle: bool,
147) -> None:
148 """Load a subimage of a cell coadd from a butler repository and save it
149 to testdata_images.
151 Parameters
152 ----------
153 butler
154 Butler repository to read the dataset from.
155 output_path
156 Path to the FITS file to write under testdata_images.
157 dataset_ref
158 Reference to the dataset to load from the butler.
159 shuffle
160 Whether to shuffle pixels within each cell to obscure proprietary
161 data.
162 """
163 from lsst.cell_coadds import MultipleCellCoadd
165 full_cell_coadd = butler.get(dataset_ref)
166 cell_coadd = MultipleCellCoadd(
167 [
168 full_cell_coadd.cells[x, y]
169 for y in range(7, 11)
170 for x in range(5, 8)
171 if {"i": y, "j": x} != DP2_COADD_MISSING_CELL
172 ],
173 grid=full_cell_coadd.grid,
174 outer_cell_size=full_cell_coadd.outer_cell_size,
175 psf_image_size=full_cell_coadd.psf_image_size,
176 common=full_cell_coadd.common,
177 )
178 if shuffle:
179 rng = np.random.default_rng()
180 for cell in cell_coadd.cells.values():
181 indices = np.arange(cell.outer.image.array.size, dtype=int)
182 rng.shuffle(indices)
183 cell.outer.image.array[:, :] = cell.outer.image.array.flat[indices].reshape(150, 150)
184 cell.outer.mask.array[:, :] = cell.outer.mask.array.flat[indices].reshape(150, 150)
185 cell.outer.variance.array[:, :] = cell.outer.variance.array.flat[indices].reshape(150, 150)
186 for n in cell.outer.noise_realizations:
187 n.array[:, :] = n.array.flat[indices].reshape(150, 150)
188 if cell.outer.mask_fractions is not None:
189 cell.outer.mask_fractions.array[:, :] = cell.outer.mask_fractions.array.flat[indices].reshape(
190 150, 150
191 )
192 os.makedirs(os.path.dirname(output_path), exist_ok=True)
193 cell_coadd.writeFits(output_path)
196def extract_camera(butler: Butler, output_path: str, dataset_ref: DatasetRef) -> None:
197 """Read camera geometry from a butler repository and save it to
198 testdata_images.
200 Parameters
201 ----------
202 butler
203 Butler repository to read the dataset from.
204 output_path
205 Path to the FITS file to write under testdata_images.
206 dataset_ref
207 Reference to the dataset to load from the butler.
208 """
209 camera = butler.get(dataset_ref)
210 os.makedirs(os.path.dirname(output_path), exist_ok=True)
211 camera.writeFits(output_path)
214def extract_skymap(butler: Butler, output_path: str, dataset_ref: DatasetRef) -> None:
215 """Read a skymap definition from a butler repository and save it to
216 testdata_images.
218 Parameters
219 ----------
220 butler
221 Butler repository to read the dataset from.
222 output_path
223 Path to the file to write under testdata_images.
224 dataset_ref
225 Reference to the dataset to load from the butler.
226 """
227 os.makedirs(os.path.dirname(output_path), exist_ok=True)
228 (path,) = butler.retrieveArtifacts(
229 [dataset_ref],
230 destination=os.path.dirname(output_path),
231 transfer="copy",
232 preserve_path=False,
233 overwrite=True,
234 )
235 if path.ospath != output_path:
236 os.rename(path.ospath, output_path)
239def find_dataset_or_raise(
240 butler: Butler, dataset_type: str, *, collections: str | None = None, **kwargs: Any
241) -> DatasetRef:
242 """Call `lsst.daf.butler.Butler.find_dataset` with the given arguments and
243 raise `LookupError` if it returns `None`.
245 Parameters
246 ----------
247 butler
248 Butler repository to search.
249 dataset_type
250 Dataset type to find.
251 collections
252 Collections to search, or `None` for the butler's defaults.
253 **kwargs
254 Data ID key-value pairs forwarded to ``find_dataset``.
255 """
256 ref = butler.find_dataset(dataset_type, collections=collections, **kwargs)
257 if ref is None:
258 raise LookupError(f"Could not find dataset {dataset_type} with data ID {kwargs}.")
259 return ref
262@click.group("extract_test_data")
263def extract_test_data() -> None:
264 """Extract test fixtures from a Rubin data repository."""
267@extract_test_data.command("dp2")
268@click.option("-b", "--butler-repo", help="Path to the butler repository.")
269@click.option("-d", "--testdata-dir", help="Path to the testdata_images directory.")
270@click.option(
271 "-c",
272 "--collection",
273 default="LSSTCam/runs/DRP/DP2",
274 help="Collection to use for most data products.",
275)
276@click.option(
277 "--visit-images/--no-visit-images",
278 default=False,
279 help="Whether to extract preliminary_visit_image or visit_image datasets.",
280)
281@click.option(
282 "--difference-images/--no-difference-images",
283 default=False,
284 help="Whether to extract difference_image datasets.",
285)
286@click.option(
287 "--deep-coadds/--no-deep-coadds",
288 default=False,
289 help="Whether to extract deep_coadd datasets.",
290)
291@click.option(
292 "--template-coadds/--no-template-coadds",
293 default=False,
294 help="Whether to extract template_coadd datasets.",
295)
296@click.option(
297 "--camera/--no-camera",
298 default=False,
299 help="Whether to extract the camera.",
300)
301@click.option(
302 "--skymap/--no-skymap",
303 default=False,
304 help="Whether to extract the skymap.",
305)
306def extract_dp2(
307 butler_repo: str | None,
308 testdata_dir: str | None,
309 collection: str,
310 *,
311 visit_images: bool,
312 difference_images: bool,
313 deep_coadds: bool,
314 template_coadds: bool,
315 camera: bool,
316 skymap: bool,
317) -> None: # numpydoc ignore=PR01
318 """Extract test data from a butler repository."""
319 try:
320 # lsst.afw.image is imported only to confirm a full Rubin pipelines
321 # environment is present: daf.butler and lsst.utils are available
322 # from the pip-installed package and so cannot gate on their own.
323 import lsst.afw.image # noqa: F401
324 from lsst.daf.butler import Butler
325 from lsst.utils import getPackageDir
326 except ImportError as err:
327 err.add_note(
328 "Updating the test data requires a full Rubin development enviroment with at least "
329 "'afw', 'obs_base', 'meas_extensions_psfex', 'meas_extensions_piff' and 'cell_coadds' "
330 "importable. This is not necessary for just running the tests."
331 )
332 raise
333 if butler_repo is None:
334 butler_repo = "dp2_prep"
335 if testdata_dir is None:
336 testdata_dir = getPackageDir("testdata_images")
337 butler = Butler.from_config(butler_repo, collections=[collection])
338 if visit_images:
339 extract_exposure(
340 butler,
341 os.path.join(testdata_dir, "dp2", "legacy", "visit_image.fits"),
342 find_dataset_or_raise(butler, "visit_image", **DP2_VISIT_DETECTOR_DATA_ID),
343 shuffle=True,
344 )
345 extract_visit_image_background(
346 butler,
347 os.path.join(testdata_dir, "dp2", "legacy", "visit_image_background.fits"),
348 find_dataset_or_raise(butler, "visit_image_background", **DP2_VISIT_DETECTOR_DATA_ID),
349 )
350 extract_generic(
351 butler,
352 os.path.join(testdata_dir, "dp2", "legacy", "visit_summary.fits"),
353 find_dataset_or_raise(butler, "visit_summary", **DP2_VISIT_DETECTOR_DATA_ID),
354 )
355 extract_exposure(
356 butler,
357 os.path.join(testdata_dir, "dp2", "legacy", "preliminary_visit_image.fits"),
358 find_dataset_or_raise(butler, "preliminary_visit_image", **DP2_VISIT_DETECTOR_DATA_ID),
359 shuffle=True,
360 )
361 extract_visit_image_background(
362 butler,
363 os.path.join(testdata_dir, "dp2", "legacy", "preliminary_visit_image_background.fits"),
364 find_dataset_or_raise(butler, "preliminary_visit_image_background", **DP2_VISIT_DETECTOR_DATA_ID),
365 )
366 if difference_images:
367 extract_exposure(
368 butler,
369 os.path.join(testdata_dir, "dp2", "legacy", "difference_image.fits"),
370 find_dataset_or_raise(butler, "difference_image", **DP2_VISIT_DETECTOR_DATA_ID),
371 shuffle=True,
372 )
373 extract_exposure(
374 butler,
375 os.path.join(testdata_dir, "dp2", "legacy", "template_detector.fits"),
376 find_dataset_or_raise(butler, "template_detector", **DP2_VISIT_DETECTOR_DATA_ID),
377 shuffle=True,
378 )
379 extract_generic(
380 butler,
381 os.path.join(testdata_dir, "dp2", "legacy", "difference_kernel.fits"),
382 find_dataset_or_raise(butler, "difference_kernel", **DP2_VISIT_DETECTOR_DATA_ID),
383 )
384 if deep_coadds:
385 extract_cell_coadd(
386 butler,
387 os.path.join(
388 testdata_dir,
389 "dp2",
390 "legacy",
391 "deep_coadd_cell_predetection.fits",
392 ),
393 find_dataset_or_raise(butler, "deep_coadd_cell_predetection", **DP2_COADD_DATA_ID),
394 shuffle=True,
395 )
396 if template_coadds:
397 extract_exposure(
398 butler,
399 os.path.join(testdata_dir, "dp2", "legacy", "template_coadd.fits"),
400 find_dataset_or_raise(butler, "template_coadd", **DP2_COADD_DATA_ID),
401 shuffle=True,
402 )
403 if skymap:
404 extract_skymap(
405 butler,
406 os.path.join(
407 testdata_dir,
408 "dp2",
409 "legacy",
410 "skyMap.pickle",
411 ),
412 find_dataset_or_raise(butler, "skyMap", skymap=DP2_COADD_DATA_ID["skymap"]),
413 )
414 if camera:
415 extract_camera(
416 butler,
417 os.path.join(
418 testdata_dir,
419 "dp2",
420 "legacy",
421 "camera.fits",
422 ),
423 find_dataset_or_raise(butler, "camera", instrument="LSSTCam"),
424 )
427if __name__ == "__main__":
428 extract_test_data()