Coverage for python/lsst/images/tests/extract_legacy_test_data.py: 20%

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. 

11 

12from __future__ import annotations 

13 

14__all__ = () 

15 

16import os 

17from typing import TYPE_CHECKING, Any 

18 

19import click 

20import numpy as np 

21 

22if TYPE_CHECKING: 

23 from lsst.daf.butler import Butler, DatasetRef 

24 

25 

26from ._data_ids import DP2_COADD_DATA_ID, DP2_COADD_MISSING_CELL, DP2_VISIT_DETECTOR_DATA_ID 

27 

28 

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 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 

58 

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 ) 

91 

92 

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. 

100 

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) 

112 

113 

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. 

121 

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) 

133 

134 

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. 

143 

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 

157 

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) 

187 

188 

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. 

192 

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) 

205 

206 

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. 

210 

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) 

230 

231 

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`. 

237 

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 

253 

254 

255@click.group("extract_test_data") 

256def extract_test_data() -> None: 

257 """Extract test fixtures from a Rubin data repository.""" 

258 

259 

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 ) 

405 

406 

407if __name__ == "__main__": 

408 extract_test_data()