Coverage for tests/test_cell_coadd.py: 40%

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 

14import dataclasses 

15import os 

16import pickle 

17from typing import Any 

18 

19import numpy as np 

20import pytest 

21 

22from lsst.images import YX, Box, Interval, MaskPlane, get_legacy_deep_coadd_mask_planes 

23from lsst.images.cells import CellCoadd, CellIJ 

24from lsst.images.fits import FitsCompressionOptions 

25from lsst.images.serialization import read_archive 

26from lsst.images.tests import ( 

27 DP2_COADD_DATA_ID, 

28 DP2_COADD_MISSING_CELL, 

29 RoundtripFits, 

30 RoundtripJson, 

31 RoundtripNdf, 

32 assert_cell_coadds_equal, 

33 assert_images_equal, 

34 assert_masked_images_equal, 

35 assert_psfs_equal, 

36 check_bounds_contains_broadcasting, 

37 compare_cell_coadd_to_legacy, 

38 compare_masked_image_to_legacy, 

39 compare_psf_to_legacy, 

40 compare_sky_projection_to_legacy_wcs, 

41) 

42 

43try: 

44 import h5py # noqa: F401 

45 

46 HAVE_H5PY = True 

47except ImportError: 

48 HAVE_H5PY = False 

49 

50try: 

51 import lsst.afw.image # noqa: F401 

52 from lsst.cell_coadds import MultipleCellCoadd as LegacyMultipleCellCoadd 

53 

54 HAVE_LEGACY = True 

55except ImportError: 

56 HAVE_LEGACY = False 

57 type LegacyMultipleCellCoadd = Any # type: ignore[no-redef] 

58 

59EXTERNAL_DATA_DIR = os.environ.get("TESTDATA_IMAGES_DIR", None) 

60LOCAL_DATA_DIR = os.path.join(os.path.dirname(__file__), "data") 

61 

62skip_no_h5py = pytest.mark.skipif(not HAVE_H5PY, reason="h5py is not installed") 

63skip_no_legacy = pytest.mark.skipif(not HAVE_LEGACY, reason="lsst.afw (etc) could not be imported.") 

64 

65 

66@dataclasses.dataclass 

67class _LegacyTestData: 

68 """A struct holding test data loaded from EXTERNAL_DATA_DIR.""" 

69 

70 filename: str 

71 tract_bbox: Box 

72 legacy_cell_coadd: LegacyMultipleCellCoadd 

73 cell_coadd: CellCoadd 

74 plane_map: dict[str, MaskPlane] = dataclasses.field(default_factory=get_legacy_deep_coadd_mask_planes) 

75 

76 def make_psf_points(self, bbox: Box | None = None) -> YX[np.ndarray]: 

77 """Create random PSF sample points within the given bbox.""" 

78 if bbox is None: 

79 bbox = self.cell_coadd.bbox 

80 rng = np.random.default_rng(44) 

81 xc, yc = np.meshgrid( 

82 np.arange( 

83 bbox.x.start + self.cell_coadd.grid.cell_shape.x * 0.5, 

84 bbox.x.stop, 

85 self.cell_coadd.grid.cell_shape.x, 

86 ), 

87 np.arange( 

88 bbox.y.start + self.cell_coadd.grid.cell_shape.y * 0.5, 

89 bbox.y.stop, 

90 self.cell_coadd.grid.cell_shape.y, 

91 ), 

92 ) 

93 return YX( 

94 y=yc.ravel() + rng.uniform(-0.4, 0.4, size=yc.size), 

95 x=xc.ravel() + rng.uniform(-0.4, 0.4, size=xc.size), 

96 ) 

97 

98 

99@pytest.fixture(scope="session") 

100def legacy_test_data() -> _LegacyTestData: 

101 """Return a struct of CellCoadd loaded from legacy test data. 

102 

103 Skips if ``TESTDATA_IMAGES_DIR`` is not set or if ``lsst.cell_coadds`` 

104 cannot be imported. 

105 """ 

106 if EXTERNAL_DATA_DIR is None: 106 ↛ 108line 106 didn't jump to line 108 because the condition on line 106 was always true

107 pytest.skip("TESTDATA_IMAGES_DIR is not in the environment.") 

108 try: 

109 from lsst.cell_coadds import MultipleCellCoadd 

110 except ImportError: 

111 pytest.skip("lsst.cell_coadds could not be imported.") 

112 filename = os.path.join(EXTERNAL_DATA_DIR, "dp2", "legacy", "deep_coadd_cell_predetection.fits") 

113 plane_map = get_legacy_deep_coadd_mask_planes() 

114 legacy_cell_coadd = MultipleCellCoadd.read_fits(filename) 

115 with open(os.path.join(EXTERNAL_DATA_DIR, "dp2", "legacy", "skyMap.pickle"), "rb") as stream: 

116 skymap = pickle.load(stream) 

117 cell_coadd = CellCoadd.from_legacy( 

118 legacy_cell_coadd, 

119 plane_map=plane_map, 

120 tract_info=skymap[DP2_COADD_DATA_ID["tract"]], 

121 ) 

122 return _LegacyTestData( 

123 filename=filename, 

124 tract_bbox=Box.from_legacy(skymap[DP2_COADD_DATA_ID["tract"]].getBBox()), 

125 legacy_cell_coadd=legacy_cell_coadd, 

126 cell_coadd=cell_coadd, 

127 ) 

128 

129 

130@pytest.fixture 

131def minified_cell_coadd() -> CellCoadd: 

132 """Return a tiny CellCoadd from JSON data stored in this package.""" 

133 path = os.path.join(LOCAL_DATA_DIR, "schema_v1", "legacy", "cell_coadd.json") 

134 return read_archive(path, CellCoadd) 

135 

136 

137def make_subbox(full_bbox: Box) -> Box: 

138 """Make a box that's useful for nontrivial subimage tests. 

139 

140 This box only overlaps (but does not fully cover) the middle 2 (of 4) 

141 cells in y, while covering exactly the last column of cells in x. It does 

142 not cover the missing cell. 

143 """ 

144 return Box.factory[ 

145 full_bbox.y.start + 252 : full_bbox.y.stop - 175, 

146 full_bbox.x.stop - 150 : full_bbox.x.stop, 

147 ] 

148 

149 

150def test_from_legacy(legacy_test_data: _LegacyTestData) -> None: 

151 """Test constructing a CellCoadd by converting a legacy 

152 ``MultipleCellCoadd``. 

153 """ 

154 assert legacy_test_data.cell_coadd.bounds.missing == {CellIJ(**DP2_COADD_MISSING_CELL)} 

155 assert legacy_test_data.cell_coadd.bbox == Box.factory[12900:13500, 9600:10050] 

156 compare_cell_coadd_to_legacy( 

157 legacy_test_data.cell_coadd, 

158 legacy_test_data.legacy_cell_coadd, 

159 tract_bbox=legacy_test_data.tract_bbox, 

160 plane_map=legacy_test_data.plane_map, 

161 psf_points=legacy_test_data.make_psf_points(), 

162 ) 

163 

164 

165def test_roundtrip(legacy_test_data: _LegacyTestData) -> None: 

166 """Test that a CellCoadd roundtrips through FITS.""" 

167 with RoundtripFits(legacy_test_data.cell_coadd, "CellCoadd") as roundtrip: 

168 # Check a subimage read (no component arg — does not trigger a skip). 

169 subbox = Box.factory[ 

170 legacy_test_data.cell_coadd.bbox.y.start + 252 : legacy_test_data.cell_coadd.bbox.y.stop - 175, 

171 legacy_test_data.cell_coadd.bbox.x.stop - 150 : legacy_test_data.cell_coadd.bbox.x.stop, 

172 ] 

173 subimage = roundtrip.get(bbox=subbox) 

174 assert_masked_images_equal(subimage, legacy_test_data.cell_coadd[subbox], expect_view=False) 

175 with roundtrip.inspect() as fits: 

176 for extname in ["IMAGE", "MASK", "VARIANCE", "MASK_FRACTIONS/REJECTED"] + [ 

177 f"NOISE_REALIZATIONS/{n}" for n in range(len(legacy_test_data.cell_coadd.noise_realizations)) 

178 ]: 

179 assert fits[extname].header["ZTILE1"] == legacy_test_data.cell_coadd.grid.cell_shape.x 

180 assert fits[extname].header["ZTILE2"] == legacy_test_data.cell_coadd.grid.cell_shape.y 

181 # Fixture self-consistency: bbox and missing-cell set are as expected. 

182 assert legacy_test_data.cell_coadd.bounds.missing == {CellIJ(**DP2_COADD_MISSING_CELL)} 

183 assert legacy_test_data.cell_coadd.bbox == Box.factory[12900:13500, 9600:10050] 

184 # Full round-trip fidelity. 

185 assert_cell_coadds_equal(roundtrip.result, legacy_test_data.cell_coadd, expect_view=False) 

186 compare_cell_coadd_to_legacy( 

187 roundtrip.result, 

188 legacy_test_data.legacy_cell_coadd, 

189 tract_bbox=legacy_test_data.tract_bbox, 

190 plane_map=legacy_test_data.plane_map, 

191 psf_points=legacy_test_data.make_psf_points(), 

192 ) 

193 

194 

195def test_roundtrip_components(legacy_test_data: _LegacyTestData) -> None: 

196 """Test component and subimage reads. 

197 

198 This test will be skipped if `lsst.daf.butler` is not available instead of 

199 falling back to non-butler I/O, which is why we don't want to merge it 

200 with `test_roundtrip`. 

201 """ 

202 with RoundtripFits(legacy_test_data.cell_coadd, "CellCoadd") as roundtrip: 

203 subbox = make_subbox(legacy_test_data.cell_coadd.bbox) 

204 subpsf = roundtrip.get("psf", bbox=subbox) 

205 assert subpsf.bounds.bbox == Box( 

206 y=Interval.factory[ 

207 legacy_test_data.cell_coadd.bbox.y.start + 150 : legacy_test_data.cell_coadd.bbox.y.stop - 150 

208 ], 

209 x=subbox.x, 

210 ) 

211 assert_psfs_equal( 

212 subpsf, 

213 legacy_test_data.cell_coadd.psf, 

214 points=legacy_test_data.make_psf_points(subbox), 

215 ) 

216 assert roundtrip.get("bbox") == legacy_test_data.cell_coadd.bbox 

217 alternates = { 

218 k: roundtrip.get(k) 

219 for k in [ 

220 "sky_projection", 

221 "image", 

222 "mask", 

223 "variance", 

224 "masked_image", 

225 "psf", 

226 "aperture_corrections", 

227 "provenance", 

228 "backgrounds", 

229 "bbox", 

230 ] 

231 } 

232 # Read all the components at once. 

233 all_components = roundtrip.get("components") 

234 assert set(all_components) == set(alternates) - {"masked_image"} 

235 assert all_components["bbox"] == alternates["bbox"] 

236 assert_psfs_equal(all_components["psf"], alternates["psf"]) 

237 assert_images_equal(all_components["image"], alternates["image"]) 

238 

239 backgrounds = roundtrip.get("backgrounds") 

240 assert backgrounds.keys() == set() 

241 assert backgrounds.subtracted is None 

242 

243 compare_cell_coadd_to_legacy( 

244 roundtrip.result, 

245 legacy_test_data.legacy_cell_coadd, 

246 tract_bbox=legacy_test_data.tract_bbox, 

247 plane_map=legacy_test_data.plane_map, 

248 alternates=alternates, 

249 psf_points=legacy_test_data.make_psf_points(), 

250 ) 

251 

252 

253def test_fits_compression(legacy_test_data: _LegacyTestData) -> None: 

254 """Test lossy FITS compression produces the expected headers.""" 

255 with RoundtripFits( 

256 legacy_test_data.cell_coadd, 

257 storage_class="CellCoadd", 

258 recipe="lossy16", 

259 compression_options={ 

260 "image": FitsCompressionOptions.LOSSY, 

261 "variance": FitsCompressionOptions.LOSSY, 

262 }, 

263 ) as roundtrip: 

264 with roundtrip.inspect() as fits: 

265 for extname in ["IMAGE", "MASK", "VARIANCE", "MASK_FRACTIONS/REJECTED"] + [ 

266 f"NOISE_REALIZATIONS/{n}" for n in range(len(legacy_test_data.cell_coadd.noise_realizations)) 

267 ]: 

268 assert fits[extname].header["ZTILE1"] == legacy_test_data.cell_coadd.grid.cell_shape.x 

269 assert fits[extname].header["ZTILE2"] == legacy_test_data.cell_coadd.grid.cell_shape.y 

270 if extname == "MASK" or extname.startswith("MASK_FRACTIONS"): 

271 assert fits[extname].header["ZCMPTYPE"] == "GZIP_2" 

272 else: 

273 assert fits[extname].header["ZCMPTYPE"] == "RICE_1" 

274 assert fits[extname].header["ZQUANTIZ"] == "SUBTRACTIVE_DITHER_2" 

275 

276 

277def test_json_roundtrip(legacy_test_data: _LegacyTestData) -> None: 

278 """Verify a CellCoadd round-trips correctly through the JSON archive.""" 

279 with RoundtripJson(legacy_test_data.cell_coadd) as roundtrip: 

280 pass 

281 assert_cell_coadds_equal(roundtrip.result, legacy_test_data.cell_coadd, expect_view=False) 

282 

283 

284def test_to_legacy_cell_coadd(legacy_test_data: _LegacyTestData) -> None: 

285 """Verify converting a CellCoadd back into a legacy MultipleCellCoadd.""" 

286 legacy_cell_coadd = legacy_test_data.cell_coadd.to_legacy_cell_coadd() 

287 compare_cell_coadd_to_legacy( 

288 legacy_test_data.cell_coadd, 

289 legacy_cell_coadd, 

290 tract_bbox=legacy_test_data.tract_bbox, 

291 plane_map=legacy_test_data.plane_map, 

292 psf_points=legacy_test_data.make_psf_points(), 

293 ) 

294 with pytest.raises( 

295 ValueError, match="MultipleCellCoadd requires its bounding box to lie on the cell grid." 

296 ): 

297 legacy_test_data.cell_coadd[make_subbox(legacy_test_data.cell_coadd.bbox)].to_legacy_cell_coadd() 

298 

299 

300@skip_no_legacy 

301def test_to_legacy(legacy_test_data: _LegacyTestData) -> None: 

302 """Test converting a CellCoadd back into a legacy Exposure.""" 

303 legacy_exposure = legacy_test_data.cell_coadd.to_legacy() 

304 assert legacy_exposure.getFilter().bandLabel == legacy_test_data.cell_coadd.band 

305 assert Box.from_legacy(legacy_exposure.getBBox()) == legacy_test_data.cell_coadd.bbox 

306 compare_masked_image_to_legacy( 

307 legacy_test_data.cell_coadd, 

308 legacy_exposure.maskedImage, 

309 plane_map=legacy_test_data.plane_map, 

310 expect_view=True, 

311 ) 

312 compare_psf_to_legacy( 

313 legacy_test_data.cell_coadd.psf, 

314 legacy_exposure.getPsf(), 

315 points=legacy_test_data.make_psf_points(), 

316 expect_legacy_raise_on_out_of_bounds=True, 

317 ) 

318 compare_sky_projection_to_legacy_wcs( 

319 legacy_test_data.cell_coadd.sky_projection, 

320 legacy_exposure.getWcs(), 

321 legacy_test_data.cell_coadd.sky_projection.pixel_frame, 

322 subimage_bbox=legacy_test_data.cell_coadd.bbox, 

323 is_fits=True, 

324 ) 

325 subbox = make_subbox(legacy_test_data.cell_coadd.bbox) 

326 compare_masked_image_to_legacy( 

327 legacy_test_data.cell_coadd[subbox], 

328 legacy_test_data.cell_coadd[subbox].to_legacy().maskedImage, 

329 plane_map=legacy_test_data.plane_map, 

330 expect_view=True, 

331 ) 

332 

333 

334@skip_no_h5py 

335def test_ndf_roundtrip(legacy_test_data: _LegacyTestData) -> None: 

336 """Test that CellCoadd round-trips through NDF.""" 

337 with RoundtripNdf(legacy_test_data.cell_coadd, "CellCoadd") as roundtrip: 

338 assert_cell_coadds_equal(roundtrip.result, legacy_test_data.cell_coadd, expect_view=False) 

339 

340 

341def test_cell_grid_bounds_contains_broadcasting(minified_cell_coadd: CellCoadd) -> None: 

342 """Test that CellGridBounds.contains broadcasts like a numpy ufunc.""" 

343 assert minified_cell_coadd.bounds.missing, "fixture should retain a missing cell" 

344 check_bounds_contains_broadcasting(minified_cell_coadd.bounds) 

345 

346 

347def test_intersection_bounds_contains_broadcasting(minified_cell_coadd: CellCoadd) -> None: 

348 """Test that IntersectionBounds.contains broadcasts like a numpy ufunc.""" 

349 # Clip the CellGridBounds with a Box offset by 1 pixel on each side so it 

350 # does not snap to any cell boundary, forcing a lazy IntersectionBounds. 

351 bounds = minified_cell_coadd.bounds 

352 clip = Box.factory[ 

353 bounds.bbox.y.start + 1 : bounds.bbox.y.stop - 1, 

354 bounds.bbox.x.start + 1 : bounds.bbox.x.stop - 1, 

355 ] 

356 check_bounds_contains_broadcasting(bounds.intersection(clip))