Coverage for python/lsst/images/tests/verify_rewrite.py: 39%

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 random 

17from collections.abc import Iterable, Sequence 

18from typing import TYPE_CHECKING, cast 

19 

20import astropy.io.fits 

21import click 

22import fsspec 

23import numpy as np 

24 

25from lsst.images import ( 

26 BackgroundMap, 

27 Box, 

28 DifferenceImage, 

29 DifferenceImageTemplateInfo, 

30 VisitImage, 

31 get_legacy_difference_image_mask_planes, 

32) 

33from lsst.images.convolution_kernels import ConvolutionKernel 

34from lsst.images.tests import compare_masked_image_to_legacy, compare_visit_image_to_legacy 

35 

36if TYPE_CHECKING: 

37 import tqdm 

38 

39 from lsst.daf.butler import Butler, DataCoordinate 

40 

41 

42# These defaults match the rewrite tasks used for DP2 in pipe_tasks. That makes 

43# this a bit of a cyclic DRY violation, but not a serious problem. 

44VISIT_IMAGE_BACKGROUNDS: tuple[str, ...] = ("subtracted", "skyCorr") 

45DIFFERENCE_IMAGE_BACKGROUNDS: tuple[str, ...] = () 

46 

47 

48def _check_kernel(kernel: ConvolutionKernel | None) -> None: 

49 """Sanity-check a DifferenceImage kernel.""" 

50 assert isinstance(kernel, ConvolutionKernel), f"kernel has type {type(kernel)}" 

51 assert kernel.bounds is not None 

52 assert kernel.bounds.bbox is not None 

53 xy = kernel.bounds.bbox.meshgrid(3) 

54 for x, y in zip(xy.x.ravel(), xy.y.ravel(), strict=True): 

55 im = kernel.compute_kernel_image(x=int(x), y=int(y)) 

56 assert im.array.size > 0 

57 assert np.isfinite(im.array).all(), "kernel has non-finite values" 

58 assert np.abs(im.array).max() > 0, "kernel is identically zero" 

59 

60 

61def _check_templates(templates: list[DifferenceImageTemplateInfo] | None, detector_bbox: Box) -> None: 

62 """Sanity-check the DifferenceImage template-info list.""" 

63 assert isinstance(templates, list), f"templates has type {type(templates)}" 

64 assert len(templates) > 0, "no template info attached" 

65 for info in templates: 

66 assert detector_bbox.contains(info.bounds.bbox), "template bounds outside detector" 

67 if info.psf_shape_flag: 

68 continue 

69 assert info.psf_shape_xx * info.psf_shape_yy - info.psf_shape_xy**2 > 0, ( 

70 "template PSF shape is not a valid ellipse" 

71 ) 

72 

73 

74def _check_backgrounds(backgrounds: BackgroundMap, bbox: Box, *, expected: Sequence[str] = ()) -> None: 

75 """Sanity-check the backgrounds attached to an image. 

76 

77 ``expected`` names must be present. Each attached background's field must 

78 cover the image bbox and evaluate to finite values over a coarse grid, and 

79 the map's ``subtracted`` attribute must point at a present background. 

80 """ 

81 present = set(backgrounds) 

82 for name in expected: 

83 assert name in present, f"expected background {name!r} not attached" 

84 for name, bg in backgrounds.items(): 

85 field = bg.field 

86 assert field.bounds.bbox.contains(bbox), f"background {name!r} does not cover the image bbox" 

87 grid = field.bounds.bbox.meshgrid(2) 

88 vals = field(x=grid.x, y=grid.y) 

89 assert np.isfinite(vals).all(), f"background {name!r} has non-finite values" 

90 if (subtracted := backgrounds.subtracted) is not None: 

91 assert subtracted.name in present, "subtracted background designation not in map" 

92 

93 

94@click.command("verify-rewrite") 

95@click.argument("repo") 

96@click.argument("dataset_type") 

97@click.argument("collection") 

98@click.option("--where", default="", help="Query string to constraint the comparison.") 

99@click.option("--old-prefix", default="legacy_", help="Prefix for the old dataset type names.") 

100@click.option("--new-prefix", default="", help="Prefix for the new dataset type names.") 

101@click.option( 

102 "--require-compressed/--no-require-compressed", 

103 default=True, 

104 help="Check that the new data product is lossy-compressed.", 

105) 

106@click.option("--check-kernel/--no-check-kernel", default=True, help="Sanity-check DifferenceImage.kernel.") 

107@click.option( 

108 "--check-templates/--no-check-templates", default=True, help="Sanity-check DifferenceImage.templates." 

109) 

110@click.option( 

111 "--check-backgrounds/--no-check-backgrounds", default=True, help="Sanity-check attached backgrounds." 

112) 

113def verify_rewrite( 

114 *, 

115 repo: str, 

116 dataset_type: str, 

117 collection: str, 

118 where: str, 

119 require_compressed: bool, 

120 old_prefix: str, 

121 new_prefix: str, 

122 check_kernel: bool, 

123 check_templates: bool, 

124 check_backgrounds: bool, 

125) -> None: # numpydoc ignore=PR01 

126 """Compare rewritten images in COLLECTION against the originals in 

127 COLLECTION of REPO. 

128 """ 

129 try: 

130 from lsst.afw.image import Exposure # noqa: F401 

131 from lsst.daf.butler import Butler 

132 except ImportError as err: 

133 err.add_note("verify-rewrite requires a full Rubin development environment.") 

134 raise 

135 with Butler.from_config(repo, collections=[collection]) as butler: 

136 verifier = RewriteVerifier(butler, dataset_type, old_prefix=old_prefix, new_prefix=new_prefix) 

137 # Query for and compare datasets. 

138 for data_id in verifier.process(where): 

139 # Test visit_image vs. future_visit_image. 

140 verifier.compare_images( 

141 data_id, 

142 check_kernel=check_kernel, 

143 check_templates=check_templates, 

144 check_backgrounds=check_backgrounds, 

145 ) 

146 if require_compressed: 

147 verifier.require_compressed(data_id) 

148 

149 if verifier.n_problems: 

150 raise click.exceptions.Exit(1) 

151 

152 

153class RewriteVerifier: 

154 def __init__( 

155 self, 

156 butler: Butler, 

157 base_dataset_type: str, 

158 *, 

159 old_prefix: str, 

160 new_prefix: str, 

161 ) -> None: 

162 self.butler = butler 

163 self.base_dataset_type = base_dataset_type 

164 self.old_prefix = old_prefix 

165 self.new_prefix = new_prefix 

166 self._progress: tqdm.tqdm | None = None 

167 self._n_problems = 0 

168 

169 def _report_problem(self) -> None: 

170 self._n_problems += 1 

171 

172 @property 

173 def n_problems(self) -> int: 

174 return self._n_problems 

175 

176 def process(self, where: str, *dimensions: str) -> Iterable[DataCoordinate]: 

177 import tqdm 

178 

179 old_dataset_type = self.butler.get_dataset_type(f"{self.old_prefix}{self.base_dataset_type}") 

180 dimension_group = old_dataset_type.dimensions.union(self.butler.dimensions.conform(dimensions)) 

181 with self.butler.query() as query: 

182 data_ids = list( 

183 tqdm.tqdm( 

184 query.where(where).join_dataset_search(old_dataset_type.name).data_ids(dimension_group), 

185 desc=f"querying for {old_dataset_type.name}", 

186 ) 

187 ) 

188 new_count = query.datasets(f"{self.new_prefix}{self.base_dataset_type}").where(where).count() 

189 assert len(data_ids) == new_count, f"Count mismatch: new ({new_count}) != old ({len(data_ids)})." 

190 random.shuffle(data_ids) 

191 self._progress = tqdm.tqdm(data_ids) 

192 yield from self._progress 

193 

194 def compare_images( 

195 self, 

196 data_id: DataCoordinate, 

197 *, 

198 check_kernel: bool = True, 

199 check_templates: bool = True, 

200 check_backgrounds: bool = True, 

201 ) -> None: 

202 old = self.butler.get(f"{self.old_prefix}{self.base_dataset_type}", data_id) 

203 new = self.butler.get(f"{self.new_prefix}{self.base_dataset_type}", data_id) 

204 expected_backgrounds = None 

205 plane_map = None 

206 if isinstance(new, DifferenceImage): 

207 expected_backgrounds = DIFFERENCE_IMAGE_BACKGROUNDS 

208 plane_map = get_legacy_difference_image_mask_planes() 

209 if check_kernel or check_templates: 

210 try: 

211 if check_kernel: 

212 _check_kernel(new.kernel) 

213 if check_templates: 

214 _check_templates(new.templates, new.bbox) 

215 except Exception as err: 

216 self.print_error(data_id, err) 

217 return 

218 if isinstance(new, VisitImage): 

219 try: 

220 compare_visit_image_to_legacy( 

221 new, 

222 old, 

223 expect_view=False, 

224 plane_map=plane_map, 

225 check_photometric_scaling=False, 

226 instrument=cast(str, data_id["instrument"]), 

227 visit=cast(int, data_id["visit"]), 

228 detector=cast(int, data_id["detector"]), 

229 ) 

230 except Exception as err: 

231 self.print_error(data_id, err) 

232 return 

233 if expected_backgrounds is None: 

234 expected_backgrounds = VISIT_IMAGE_BACKGROUNDS 

235 else: 

236 try: 

237 compare_masked_image_to_legacy(new, old, expect_view=False, plane_map=plane_map) 

238 except Exception as err: 

239 self.print_error(data_id, err) 

240 return 

241 if check_backgrounds and expected_backgrounds is not None: 

242 _check_backgrounds(new.backgrounds, new.bbox, expected=expected_backgrounds) 

243 

244 def require_compressed(self, data_id: DataCoordinate) -> None: 

245 ref = self.butler.find_dataset(f"{self.new_prefix}{self.base_dataset_type}", data_id) 

246 assert ref is not None, f"new dataset for {data_id} is missing" 

247 path = self.butler.getURI(ref) 

248 fs: fsspec.AbstractFileSystem 

249 fs, fp = path.to_fsspec() 

250 with fs.open(fp) as stream: 

251 with astropy.io.fits.open(stream, disable_image_compression=True) as hdu_list: 

252 if (image_zcmptype := hdu_list["IMAGE"].header.get("ZCMPTYPE")) != "RICE_1": 

253 self._report_problem() 

254 self.print(data_id, f"IMAGE HDU has ZCMPTYPE={image_zcmptype!r}") 

255 if (mask_zcmptype := hdu_list["MASK"].header.get("ZCMPTYPE")) != "GZIP_2": 

256 self._report_problem() 

257 self.print(data_id, f"MASK HDU has ZCMPTYPE={mask_zcmptype!r}") 

258 if (variance_zcmptype := hdu_list["VARIANCE"].header.get("ZCMPTYPE")) != "RICE_1": 

259 self._report_problem() 

260 self.print(data_id, f"VARIANCE HDU has ZCMPTYPE={variance_zcmptype!r}") 

261 

262 def print(self, data_id: DataCoordinate, msg: str) -> None: 

263 if self._progress is not None: 

264 self._progress.write(f"{data_id}: {msg}") 

265 else: 

266 print(f"{data_id}: {msg}") 

267 

268 def print_error(self, data_id: DataCoordinate, err: Exception) -> None: 

269 self._report_problem() 

270 message = f"{err} ({type(err).__name__})" 

271 notes = getattr(err, "__notes__", ()) 

272 if notes: 

273 # The comparison utilities attach the failing component's name as 

274 # an exception note (a hierarchy from outer to inner when blocks 

275 # are nested, though they never are today). Prefix the message 

276 # with that path so it reads as 

277 # "<data-id>\n <component>: <message>" 

278 # rather than an indented note trailing the message. 

279 component = " -> ".join(notes) 

280 block = f"{data_id}\n {component}: {message}" 

281 else: 

282 block = f"{data_id}: {message}" 

283 if self._progress is not None: 283 ↛ 284line 283 didn't jump to line 284 because the condition on line 283 was never true

284 self._progress.write(block) 

285 else: 

286 print(block) 

287 

288 

289if __name__ == "__main__": 

290 verify_rewrite()