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