Coverage for tests/test_masked_image.py: 91%

208 statements  

« prev     ^ index     » next       coverage.py v7.15.2, created at 2026-07-23 07:45 +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. 

11 

12from __future__ import annotations 

13 

14import dataclasses 

15import os 

16from pathlib import Path 

17from typing import Any 

18 

19import astropy.io.fits 

20import astropy.units as u 

21import numpy as np 

22import pytest 

23from astropy.coordinates import Angle, SkyCoord 

24 

25from lsst.images import ( 

26 Box, 

27 GeneralFrame, 

28 Image, 

29 MaskedImage, 

30 MaskPlane, 

31 MaskSchema, 

32 NotContainedError, 

33 SkyProjection, 

34 get_legacy_visit_image_mask_planes, 

35) 

36from lsst.images.fits import FitsCompressionOptions 

37from lsst.images.tests import ( 

38 RoundtripFits, 

39 RoundtripJson, 

40 RoundtripNdf, 

41 assert_masked_images_equal, 

42 compare_masked_image_to_legacy, 

43) 

44 

45try: 

46 import h5py # noqa: F401 

47 

48 HAVE_H5PY = True 

49except ImportError: 

50 HAVE_H5PY = False 

51 

52try: 

53 from lsst.afw.image import MaskedImageReader as LegacyMaskedImageReader 

54 

55except ImportError: 

56 type LegacyMaskedImageReader = Any # type: ignore[no-redef] 

57 

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

59 

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

61 

62 

63@dataclasses.dataclass 

64class _LegacyTestData: 

65 masked_image: MaskedImage 

66 reader: LegacyMaskedImageReader 

67 plane_map: dict[str, MaskPlane] 

68 

69 

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

71def legacy_test_data() -> _LegacyTestData: 

72 """Return a Mask read directly from the legacy test dataset and a legacy 

73 reader for that image. 

74 

75 Skips if TESTDATA_IMAGES_DIR is unset or lsst.afw.image is unavailable. 

76 """ 

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

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

79 try: 

80 from lsst.afw.image import MaskedImageFitsReader 

81 except ImportError: 

82 pytest.skip("'lsst.afw.image' could not be imported.") 

83 filename = os.path.join(EXTERNAL_DATA_DIR, "dp2", "legacy", "visit_image.fits") 

84 plane_map = get_legacy_visit_image_mask_planes() 

85 masked_image = MaskedImage.read_legacy(filename, plane_map=plane_map) 

86 reader = MaskedImageFitsReader(filename) 

87 return _LegacyTestData(masked_image=masked_image, reader=reader, plane_map=plane_map) 

88 

89 

90def _make_wcs() -> astropy.wcs.WCS: 

91 """Build a gnomonic FITS WCS with 0.1 arcsec pixels at (12, 13) deg. 

92 

93 The reference pixel is at 0-based pixel (x=5, y=6). 

94 """ 

95 wcs = astropy.wcs.WCS(naxis=2) 

96 # FITS CRPIX is 1-based, so CRPIX (6, 7) is 0-based pixel (x=5, y=6). 

97 wcs.wcs.crpix = [6.0, 7.0] 

98 wcs.wcs.crval = [12.0, 13.0] 

99 scale = 0.1 / 3600.0 

100 wcs.wcs.cd = [[-scale, 0.0], [0.0, scale]] 

101 wcs.wcs.ctype = ["RA---TAN", "DEC--TAN"] 

102 return wcs 

103 

104 

105def make_masked_image() -> MaskedImage: 

106 """Return a freshly-constructed MaskedImage with BAD and HUNGRY mask 

107 planes set. 

108 """ 

109 rng = np.random.default_rng(500) 

110 masked_image = MaskedImage( 

111 Image(rng.normal(100.0, 8.0, size=(200, 251)), dtype=np.float64, unit=u.nJy, yx0=(5, 8)), 

112 mask_schema=MaskSchema( 

113 [ 

114 MaskPlane("BAD", "Pixel is very bad, possibly downright evil."), 

115 MaskPlane("HUNGRY", "Pixel hasn't had enough to eat today."), 

116 ] 

117 ), 

118 metadata={"fifty": "5 * 10"}, 

119 sky_projection=SkyProjection.from_fits_wcs(_make_wcs(), GeneralFrame(unit=u.pix)), 

120 ) 

121 masked_image.mask.array |= np.multiply.outer( 

122 masked_image.image.array < 102.0, 

123 masked_image.mask.schema.bitmask("BAD"), 

124 ) 

125 masked_image.mask.array |= np.multiply.outer( 

126 masked_image.image.array > 98.0, 

127 masked_image.mask.schema.bitmask("HUNGRY"), 

128 ) 

129 masked_image.variance.array = rng.normal(64.0, 0.5, size=masked_image.bbox.shape) 

130 return masked_image 

131 

132 

133def test_construction() -> None: 

134 """Verify the MaskedImage constructed by make_masked_image has the 

135 expected attributes. 

136 """ 

137 mi = make_masked_image() 

138 assert mi.bbox == Box.factory[5:205, 8:259] 

139 assert mi.mask.bbox == mi.bbox 

140 assert mi.variance.bbox == mi.bbox 

141 assert mi.image.array.shape == mi.bbox.shape 

142 assert mi.mask.array.shape == mi.bbox.shape + (1,) 

143 assert mi.variance.array.shape == mi.bbox.shape 

144 assert mi.unit == u.nJy 

145 assert mi.variance.unit == u.nJy**2 

146 assert mi.metadata == {"fifty": "5 * 10"} 

147 # The checks below are subject to the vagaries of the RNG, but we want 

148 # the seed to be such that they all pass, or other tests will be weaker. 

149 assert np.sum(mi.mask.array == mi.mask.schema.bitmask("BAD")) > 0 

150 assert np.sum(mi.mask.array == mi.mask.schema.bitmask("HUNGRY")) > 0 

151 assert np.sum(mi.mask.array == mi.mask.schema.bitmask("BAD", "HUNGRY")) > 0 

152 

153 assert mi[...] is not mi 

154 assert str(mi) == "MaskedImage(Image([y=5:205, x=8:259], float64), ['BAD', 'HUNGRY'])" 

155 assert ( 

156 repr(mi) 

157 == "MaskedImage(Image(..., bbox=Box(y=Interval(start=5, stop=205), x=Interval(start=8, stop=259)), " 

158 "dtype=dtype('float64')), mask_schema=MaskSchema([MaskPlane(name='BAD', description='Pixel is " 

159 "very bad, possibly downright evil.'), MaskPlane(name='HUNGRY', description=\"Pixel hasn't had " 

160 "enough to eat today.\")], dtype=dtype('uint8')))" 

161 ) 

162 copy = mi.copy() 

163 original = mi.image.array[0, 0] 

164 copy.image.array[0, 0] = 38.0 

165 assert mi.image.array[0, 0] == original 

166 assert copy.image.array[0, 0] == 38.0 

167 

168 # Test error conditions. 

169 with pytest.raises(ValueError): 

170 # Disagreement over mask bbox. 

171 MaskedImage(Image(42.0, shape=(5, 6)), mask=mi.mask) 

172 with pytest.raises(TypeError): 

173 # No mask definition. 

174 MaskedImage(mi.image, variance=mi.variance) 

175 with pytest.raises(TypeError): 

176 # Can not provide mask and mask schema. 

177 MaskedImage( 

178 Image(42.0, shape=(5, 5)), 

179 mask=mi.mask, 

180 mask_schema=mi.mask.schema, 

181 ) 

182 with pytest.raises(ValueError): 

183 # image and variance bbox disagreement. 

184 MaskedImage( 

185 Image(42.0, shape=(5, 5)), 

186 mask_schema=mi.mask.schema, 

187 variance=mi.variance, 

188 ) 

189 with pytest.raises(ValueError): 

190 # no image unit but there is variance unit. 

191 MaskedImage( 

192 Image(42.0, shape=(5, 5)), 

193 mask_schema=mi.mask.schema, 

194 variance=Image(1.0, shape=(5, 5), unit=u.nJy), 

195 ) 

196 with pytest.raises(ValueError): 

197 # image and variance units disagree. 

198 MaskedImage( 

199 Image(42.0, shape=(5, 5), unit=u.nJy), 

200 mask_schema=mi.mask.schema, 

201 variance=Image(1.0, shape=(5, 5), unit=u.nJy), 

202 ) 

203 

204 

205def test_subset() -> None: 

206 """Verify assignment of a subset into a MaskedImage copy.""" 

207 mi = make_masked_image() 

208 copy = mi.copy() 

209 subset = copy.local[0:10, 20:30].copy() 

210 subset.image[...] = Image(42.0, shape=(10, 10), unit=u.nJy) 

211 copy[subset.bbox] = subset 

212 assert copy.image.array[0, 20] == 42.0 

213 assert copy.image.array[0, 0] == mi.image.array[0, 0] 

214 

215 

216def test_mask_setter() -> None: 

217 """Verify the mask plane can be replaced with one grown by add_plane.""" 

218 mi = make_masked_image() 

219 bad = mi.mask.get("BAD") 

220 mi.mask = mi.mask.add_plane("OUTSIDE_STENCIL", "Pixel lies outside the stencil.") 

221 assert "OUTSIDE_STENCIL" in mi.mask.schema.names 

222 assert mi.mask.bbox == mi.image.bbox 

223 np.testing.assert_array_equal(mi.mask.get("BAD"), bad) 

224 assert not mi.mask.get("OUTSIDE_STENCIL").any() 

225 # A mask whose bounding box disagrees with the image is rejected. 

226 with pytest.raises(ValueError): 

227 mi.mask = mi.mask[Box.factory[10:20, 12:22]] 

228 

229 

230def test_fits_roundtrip() -> None: 

231 """Verify MaskedImage round-trips correctly through FITS, including 

232 subimage reads. 

233 """ 

234 mi = make_masked_image() 

235 subbox = Box.factory[11:20, 25:30] 

236 subslices = (slice(6, 15), slice(17, 22)) 

237 np.testing.assert_array_equal(mi.image.array[subslices], mi.image[subbox].array) 

238 with RoundtripFits(mi, "MaskedImageV2") as roundtrip: 

239 subimage = roundtrip.get(bbox=subbox) 

240 # Check that we used lossless compression (the default). 

241 fits = roundtrip.inspect() 

242 assert fits[1].header["ZCMPTYPE"] == "GZIP_2" 

243 assert fits[2].header["ZCMPTYPE"] == "GZIP_2" 

244 assert fits[3].header["ZCMPTYPE"] == "GZIP_2" 

245 assert_masked_images_equal(roundtrip.result, mi, expect_view=False) 

246 assert_masked_images_equal(subimage, roundtrip.result[subbox], expect_view=False) 

247 

248 

249def test_fits_roundtrip_legacy_read() -> None: 

250 """Verify a round-tripped MaskedImageV2 can be read back as a legacy afw 

251 MaskedImage. 

252 """ 

253 try: 

254 import lsst.afw.image 

255 except ImportError: 

256 pytest.skip("afw could not be imported") 

257 mi = make_masked_image() 

258 with RoundtripFits(mi, "MaskedImageV2") as roundtrip: 

259 legacy_masked_image = roundtrip.get(storageClass="MaskedImage") 

260 assert isinstance(legacy_masked_image, lsst.afw.image.MaskedImage) 

261 compare_masked_image_to_legacy(mi, legacy_masked_image, expect_view=False) 

262 

263 

264def test_fits_roundtrip_lossy(tmp_path: Path) -> None: 

265 """Verify MaskedImage round-trips correctly through FITS with lossy 

266 compression. 

267 """ 

268 mi = make_masked_image() 

269 subbox = Box.factory[11:20, 25:30] 

270 subslices = (slice(6, 15), slice(17, 22)) 

271 np.testing.assert_array_equal(mi.image.array[subslices], mi.image[subbox].array) 

272 path = tmp_path / "lossy.fits" 

273 mi.write( 

274 path, 

275 compression_options={ 

276 "image": FitsCompressionOptions.LOSSY, 

277 "variance": FitsCompressionOptions.LOSSY, 

278 }, 

279 compression_seed=50, 

280 ) 

281 roundtripped = MaskedImage.read(path) 

282 subimage = MaskedImage.read(path, bbox=subbox) 

283 with astropy.io.fits.open(path, disable_image_compression=True) as fits: 

284 assert fits[1].header["ZCMPTYPE"] == "RICE_1" 

285 assert fits[2].header["ZCMPTYPE"] == "GZIP_2" 

286 assert fits[3].header["ZCMPTYPE"] == "RICE_1" 

287 assert_masked_images_equal(roundtripped, mi, expect_view=False, rtol=0.01) 

288 assert_masked_images_equal(subimage, roundtripped[subbox], expect_view=False) 

289 

290 

291@skip_no_h5py 

292def test_round_trip_ndf_compatible_mask() -> None: 

293 """Verify NDF round-trip for a MaskedImage with ≤8 mask planes.""" 

294 mi = make_masked_image() 

295 with RoundtripNdf(mi, "MaskedImageV2") as roundtrip: 

296 assert_masked_images_equal(roundtrip.result, mi, expect_view=False) 

297 

298 

299@skip_no_h5py 

300def test_round_trip_ndf_incompatible_mask() -> None: 

301 """Verify NDF round-trip for a MaskedImage with more than 8 mask planes.""" 

302 rng = np.random.default_rng(7) 

303 planes = [MaskPlane(f"P{i}", f"plane {i}") for i in range(12)] 

304 wide = MaskedImage( 

305 Image( 

306 rng.normal(100.0, 8.0, size=(50, 60)), 

307 dtype=np.float64, 

308 unit=u.nJy, 

309 yx0=(0, 0), 

310 ), 

311 mask_schema=MaskSchema(planes), 

312 ) 

313 wide.variance.array = rng.normal(64.0, 0.5, size=wide.bbox.shape) 

314 with RoundtripNdf(wide, "MaskedImageV2") as roundtrip: 

315 assert_masked_images_equal(roundtrip.result, wide, expect_view=False) 

316 

317 

318@skip_no_h5py 

319def test_round_trip_ndf_many_plane_mask() -> None: 

320 """Verify NDF round-trip for a mask that needs more than one int32 

321 chunk. 

322 """ 

323 rng = np.random.default_rng(11) 

324 planes = [MaskPlane(f"P{i}", f"plane {i}") for i in range(40)] 

325 wide = MaskedImage( 

326 Image( 

327 rng.normal(100.0, 8.0, size=(10, 12)), 

328 dtype=np.float64, 

329 unit=u.nJy, 

330 yx0=(0, 0), 

331 ), 

332 mask_schema=MaskSchema(planes), 

333 ) 

334 wide.mask.set("P0", wide.image.array > 100.0) 

335 wide.mask.set("P17", wide.image.array < 95.0) 

336 wide.mask.set("P39", wide.image.array > 110.0) 

337 wide.variance.array = rng.normal(64.0, 0.5, size=wide.bbox.shape) 

338 with RoundtripNdf(wide, "MaskedImageV2") as roundtrip: 

339 assert_masked_images_equal(roundtrip.result, wide, expect_view=False) 

340 

341 

342@skip_no_h5py 

343def test_fits_ndf_consistency() -> None: 

344 """Verify FITS and NDF backends produce equal MaskedImages on 

345 round-trip. 

346 """ 

347 mi = make_masked_image() 

348 with ( 

349 RoundtripFits(mi) as fits_rt, 

350 RoundtripNdf(mi) as ndf_rt, 

351 ): 

352 assert_masked_images_equal(mi, fits_rt.result, expect_view=False) 

353 assert_masked_images_equal(mi, ndf_rt.result, expect_view=False) 

354 assert_masked_images_equal(fits_rt.result, ndf_rt.result, expect_view=False) 

355 

356 

357def test_fits_json_consistency() -> None: 

358 """Verify FITS and JSON backends produce equal MaskedImages on 

359 round-trip. 

360 """ 

361 mi = make_masked_image() 

362 with ( 

363 RoundtripFits(mi) as fits_rt, 

364 RoundtripJson(mi) as json_rt, 

365 ): 

366 assert_masked_images_equal(mi, fits_rt.result, expect_view=False) 

367 assert_masked_images_equal(mi, json_rt.result, expect_view=False) 

368 assert_masked_images_equal(fits_rt.result, json_rt.result, expect_view=False) 

369 

370 

371def test_legacy(legacy_test_data: _LegacyTestData) -> None: 

372 """Test MaskedImage.read_legacy, MaskedImage.to_legacy, and 

373 MaskedImage.from_legacy. 

374 """ 

375 legacy_masked_image = legacy_test_data.reader.read() 

376 compare_masked_image_to_legacy( 

377 legacy_test_data.masked_image, 

378 legacy_masked_image, 

379 plane_map=legacy_test_data.plane_map, 

380 expect_view=False, 

381 ) 

382 compare_masked_image_to_legacy( 

383 legacy_test_data.masked_image, 

384 legacy_test_data.masked_image.to_legacy(plane_map=legacy_test_data.plane_map), 

385 plane_map=legacy_test_data.plane_map, 

386 expect_view=True, 

387 ) 

388 compare_masked_image_to_legacy( 

389 MaskedImage.from_legacy(legacy_masked_image, plane_map=legacy_test_data.plane_map), 

390 legacy_masked_image, 

391 expect_view=True, 

392 plane_map=legacy_test_data.plane_map, 

393 ) 

394 

395 

396def test_sky_circle_bbox() -> None: 

397 """Test that we can extract a bounding box from a sky circle.""" 

398 mi = make_masked_image() 

399 

400 # This position is on the reference pixel (x=5, y=6), which is just 

401 # outside the image (the bbox starts at x=8, y=5), so the box must be 

402 # clipped on the low-x and low-y sides. 0.1 arcsec pixels. 

403 bbox = mi.bbox_from_sky_circle( 

404 SkyCoord(ra=12.0 * u.deg, dec=13.0 * u.deg, frame="icrs"), Angle(1.0 * u.arcsec), clip=True 

405 ) 

406 # The circle has a ~10 pixel radius (1 arcsec at 0.1 arcsec per pixel), 

407 # spanning x [-5, 15] and y [-4, 16] before clipping to the image bounds. 

408 assert bbox == Box.factory[5:17, 8:16] 

409 

410 with pytest.raises(NotContainedError): 

411 # Partially off the edge but clipping not requested. 

412 mi.bbox_from_sky_circle( 

413 SkyCoord(ra=12.0 * u.deg, dec=13.0 * u.deg, frame="icrs"), Angle(1.0 * u.arcsec) 

414 ) 

415 

416 # Fully inside the image. The image is only ~200 pixels across or 

417 # ~20 arcsec. 

418 bbox = mi.bbox_from_sky_circle( 

419 SkyCoord(ra=12.0 * u.deg - 5.0 * u.arcsec, dec=13.0 * u.deg + 5.0 * u.arcsec, frame="icrs"), 

420 Angle(1.0 * u.arcsec), 

421 ) 

422 # The center is offset from the reference pixel by +50 pixels in y and 

423 # by +48.7 pixels in x (the 5 arcsec RA offset scales by cos(dec)), 

424 # placing it at (x=53.7, y=56) with a ~10 pixel radius. 

425 assert bbox == Box.factory[46:67, 44:65] 

426 

427 # Fully off the image. 

428 with pytest.raises(NotContainedError): 

429 mi.bbox_from_sky_circle( 

430 SkyCoord(ra=13.0 * u.deg, dec=13.0 * u.deg, frame="icrs"), Angle(1.0 * u.arcsec) 

431 ) 

432 

433 # Fully off the image with clipping requested. 

434 with pytest.raises(NotContainedError): 

435 mi.bbox_from_sky_circle( 

436 SkyCoord(ra=13.0 * u.deg, dec=13.0 * u.deg, frame="icrs"), Angle(1.0 * u.arcsec), clip=True 

437 ) 

438 

439 # Non-scalar center and radius are rejected. 

440 with pytest.raises(ValueError, match="scalar SkyCoord"): 

441 mi.bbox_from_sky_circle( 

442 SkyCoord(ra=[12.0, 12.1] * u.deg, dec=[13.0, 13.1] * u.deg, frame="icrs"), 

443 Angle(1.0 * u.arcsec), 

444 ) 

445 with pytest.raises(ValueError, match="scalar Angle"): 

446 mi.bbox_from_sky_circle( 

447 SkyCoord(ra=12.0 * u.deg, dec=13.0 * u.deg, frame="icrs"), Angle([1.0, 2.0] * u.arcsec) 

448 ) 

449 

450 # An image without a sky projection cannot calculate a bounding box. 

451 no_wcs = Image(0.0, shape=(10, 10), dtype=np.float64) 

452 with pytest.raises(ValueError, match="sky projection"): 

453 no_wcs.bbox_from_sky_circle( 

454 SkyCoord(ra=12.0 * u.deg, dec=13.0 * u.deg, frame="icrs"), Angle(1.0 * u.arcsec) 

455 )