Coverage for python/lsst/images/tests/extract_legacy_test_data.py: 21%
98 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-27 01:38 -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_visit_summary(
115 butler: Butler,
116 output_path: str,
117 dataset_ref: DatasetRef,
118) -> None:
119 """Load a visit_summary dataset and save it to testdata_images.
121 Parameters
122 ----------
123 butler
124 Butler repository to read the dataset from.
125 output_path
126 Path to the FITS file to write under testdata_images.
127 dataset_ref
128 Reference to the dataset to load from the butler.
129 """
130 visit_summary = butler.get(dataset_ref)
131 visit_summary.writeFits(output_path)
134def extract_cell_coadd(
135 butler: Butler,
136 output_path: str,
137 dataset_ref: DatasetRef,
138 shuffle: bool,
139) -> None:
140 """Load a subimage of a cell coadd from a butler repository and save it
141 to testdata_images.
143 Parameters
144 ----------
145 butler
146 Butler repository to read the dataset from.
147 output_path
148 Path to the FITS file to write under testdata_images.
149 dataset_ref
150 Reference to the dataset to load from the butler.
151 shuffle
152 Whether to shuffle pixels within each cell to obscure proprietary
153 data.
154 """
155 from lsst.cell_coadds import MultipleCellCoadd
157 full_cell_coadd = butler.get(dataset_ref)
158 cell_coadd = MultipleCellCoadd(
159 [
160 full_cell_coadd.cells[x, y]
161 for y in range(7, 11)
162 for x in range(5, 8)
163 if {"i": y, "j": x} != DP2_COADD_MISSING_CELL
164 ],
165 grid=full_cell_coadd.grid,
166 outer_cell_size=full_cell_coadd.outer_cell_size,
167 psf_image_size=full_cell_coadd.psf_image_size,
168 common=full_cell_coadd.common,
169 )
170 if shuffle:
171 rng = np.random.default_rng()
172 for cell in cell_coadd.cells.values():
173 indices = np.arange(cell.outer.image.array.size, dtype=int)
174 rng.shuffle(indices)
175 cell.outer.image.array[:, :] = cell.outer.image.array.flat[indices].reshape(150, 150)
176 cell.outer.mask.array[:, :] = cell.outer.mask.array.flat[indices].reshape(150, 150)
177 cell.outer.variance.array[:, :] = cell.outer.variance.array.flat[indices].reshape(150, 150)
178 for n in cell.outer.noise_realizations:
179 n.array[:, :] = n.array.flat[indices].reshape(150, 150)
180 if cell.outer.mask_fractions is not None:
181 cell.outer.mask_fractions.array[:, :] = cell.outer.mask_fractions.array.flat[indices].reshape(
182 150, 150
183 )
184 os.makedirs(os.path.dirname(output_path), exist_ok=True)
185 cell_coadd.writeFits(output_path)
188def extract_camera(butler: Butler, output_path: str, dataset_ref: DatasetRef) -> None:
189 """Read camera geometry from a butler repository and save it to
190 testdata_images.
192 Parameters
193 ----------
194 butler
195 Butler repository to read the dataset from.
196 output_path
197 Path to the FITS file to write under testdata_images.
198 dataset_ref
199 Reference to the dataset to load from the butler.
200 """
201 camera = butler.get(dataset_ref)
202 os.makedirs(os.path.dirname(output_path), exist_ok=True)
203 camera.writeFits(output_path)
206def extract_skymap(butler: Butler, output_path: str, dataset_ref: DatasetRef) -> None:
207 """Read a skymap definition from a butler repository and save it to
208 testdata_images.
210 Parameters
211 ----------
212 butler
213 Butler repository to read the dataset from.
214 output_path
215 Path to the file to write under testdata_images.
216 dataset_ref
217 Reference to the dataset to load from the butler.
218 """
219 os.makedirs(os.path.dirname(output_path), exist_ok=True)
220 (path,) = butler.retrieveArtifacts(
221 [dataset_ref],
222 destination=os.path.dirname(output_path),
223 transfer="copy",
224 preserve_path=False,
225 overwrite=True,
226 )
227 if path.ospath != output_path:
228 os.rename(path.ospath, output_path)
231def find_dataset_or_raise(
232 butler: Butler, dataset_type: str, *, collections: str | None = None, **kwargs: Any
233) -> DatasetRef:
234 """Call `lsst.daf.butler.Butler.find_dataset` with the given arguments and
235 raise `LookupError` if it returns `None`.
237 Parameters
238 ----------
239 butler
240 Butler repository to search.
241 dataset_type
242 Dataset type to find.
243 collections
244 Collections to search, or `None` for the butler's defaults.
245 **kwargs
246 Data ID key-value pairs forwarded to ``find_dataset``.
247 """
248 ref = butler.find_dataset(dataset_type, collections=collections, **kwargs)
249 if ref is None:
250 raise LookupError(f"Could not find dataset {dataset_type} with data ID {kwargs}.")
251 return ref
254@click.group("extract_test_data")
255def extract_test_data() -> None:
256 """Extract test fixtures from a Rubin data repository."""
259@extract_test_data.command("dp2")
260@click.option("-b", "--butler-repo", help="Path to the butler repository.")
261@click.option("-d", "--testdata-dir", help="Path to the testdata_images directory.")
262@click.option(
263 "-c",
264 "--collection",
265 default="LSSTCam/runs/DRP/DP2",
266 help="Collection to use for most data products.",
267)
268@click.option(
269 "--visit-images/--no-visit-images",
270 default=True,
271 help="Whether to extract preliminary_visit_image or visit_image datasets.",
272)
273@click.option(
274 "--difference-images/--no-difference-images",
275 default=True,
276 help="Whether to extract difference_image datasets.",
277)
278@click.option(
279 "--coadds/--no-coadds",
280 default=True,
281 help="Whether to extract coadd datasets.",
282)
283@click.option(
284 "--camera/--no-camera",
285 default=True,
286 help="Whether to extract the camera.",
287)
288@click.option(
289 "--skymap/--no-skymap",
290 default=True,
291 help="Whether to extract the skymap.",
292)
293def extract_dp2(
294 butler_repo: str | None,
295 testdata_dir: str | None,
296 collection: str,
297 *,
298 visit_images: bool,
299 difference_images: bool,
300 coadds: bool,
301 camera: bool,
302 skymap: bool,
303) -> None: # numpydoc ignore=PR01
304 """Extract test data from a butler repository."""
305 try:
306 # lsst.afw.image is imported only to confirm a full Rubin pipelines
307 # environment is present: daf.butler and lsst.utils are available
308 # from the pip-installed package and so cannot gate on their own.
309 import lsst.afw.image # noqa: F401
310 from lsst.daf.butler import Butler
311 from lsst.utils import getPackageDir
312 except ImportError as err:
313 err.add_note(
314 "Updating the test data requires a full Rubin development enviroment with at least "
315 "'afw', 'obs_base', 'meas_extensions_psfex', 'meas_extensions_piff' and 'cell_coadds' "
316 "importable. This is not necessary for just running the tests."
317 )
318 raise
319 if butler_repo is None:
320 butler_repo = "dp2_prep"
321 if testdata_dir is None:
322 testdata_dir = getPackageDir("testdata_images")
323 butler = Butler.from_config(butler_repo, collections=[collection])
324 if visit_images:
325 extract_exposure(
326 butler,
327 os.path.join(testdata_dir, "dp2", "legacy", "visit_image.fits"),
328 find_dataset_or_raise(butler, "visit_image", **DP2_VISIT_DETECTOR_DATA_ID),
329 shuffle=True,
330 )
331 extract_visit_image_background(
332 butler,
333 os.path.join(testdata_dir, "dp2", "legacy", "visit_image_background.fits"),
334 find_dataset_or_raise(butler, "visit_image_background", **DP2_VISIT_DETECTOR_DATA_ID),
335 )
336 extract_visit_summary(
337 butler,
338 os.path.join(testdata_dir, "dp2", "legacy", "visit_summary.fits"),
339 find_dataset_or_raise(butler, "visit_summary", **DP2_VISIT_DETECTOR_DATA_ID),
340 )
341 extract_exposure(
342 butler,
343 os.path.join(testdata_dir, "dp2", "legacy", "preliminary_visit_image.fits"),
344 find_dataset_or_raise(butler, "preliminary_visit_image", **DP2_VISIT_DETECTOR_DATA_ID),
345 shuffle=True,
346 )
347 extract_visit_image_background(
348 butler,
349 os.path.join(testdata_dir, "dp2", "legacy", "preliminary_visit_image_background.fits"),
350 find_dataset_or_raise(butler, "preliminary_visit_image_background", **DP2_VISIT_DETECTOR_DATA_ID),
351 )
352 if difference_images:
353 extract_exposure(
354 butler,
355 os.path.join(testdata_dir, "dp2", "legacy", "difference_image.fits"),
356 find_dataset_or_raise(butler, "difference_image", **DP2_VISIT_DETECTOR_DATA_ID),
357 shuffle=True,
358 )
359 if coadds:
360 extract_cell_coadd(
361 butler,
362 os.path.join(
363 testdata_dir,
364 "dp2",
365 "legacy",
366 "deep_coadd_cell_predetection.fits",
367 ),
368 find_dataset_or_raise(butler, "deep_coadd_cell_predetection", **DP2_COADD_DATA_ID),
369 shuffle=True,
370 )
371 if skymap:
372 extract_skymap(
373 butler,
374 os.path.join(
375 testdata_dir,
376 "dp2",
377 "legacy",
378 "skyMap.pickle",
379 ),
380 find_dataset_or_raise(butler, "skyMap", skymap=DP2_COADD_DATA_ID["skymap"]),
381 )
382 if camera:
383 extract_camera(
384 butler,
385 os.path.join(
386 testdata_dir,
387 "dp2",
388 "legacy",
389 "camera.fits",
390 ),
391 find_dataset_or_raise(butler, "camera", instrument="LSSTCam"),
392 )
395if __name__ == "__main__":
396 extract_test_data()