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

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 difflib 

17import random 

18from collections.abc import Iterable, Set 

19from typing import TYPE_CHECKING, Any 

20 

21import astropy.io.fits 

22import click 

23import fsspec 

24import numpy as np 

25 

26if TYPE_CHECKING: 

27 import tqdm 

28 

29 from lsst.daf.butler import Butler, DataCoordinate 

30 

31 

32@click.group("verify-rewrite") 

33def verify_rewrite() -> None: 

34 """Compare rewritten datasets against their originals in a butler repo.""" 

35 

36 

37@verify_rewrite.command("stage4") 

38@click.argument("repo") 

39@click.argument("old_collection") 

40@click.argument("new_collection") 

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

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

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

44def verify_stage4_rewrite( 

45 *, repo: str, old_collection: str, new_collection: str, where: str, old_prefix: str, new_prefix: str 

46) -> None: 

47 """Compare rewritten visit/difference images and downstream tables in 

48 NEW_COLLECTION against the originals in OLD_COLLECTION of REPO. 

49 """ 

50 try: 

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

52 from lsst.daf.butler import Butler 

53 except ImportError as err: 

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

55 raise 

56 with Butler.from_config(repo) as butler: 

57 verifier = RewriteVerifier( 

58 butler, 

59 old_collection=old_collection, 

60 new_collection=new_collection, 

61 old_prefix=old_prefix, 

62 new_prefix=new_prefix, 

63 ) 

64 # Query for and compare datasets with {visit, detector} data IDs. 

65 for data_id in verifier.process("visit_image", where, renamed=True): 

66 # Test visit_image vs. future_visit_image. 

67 verifier.compare_images("visit_image", data_id, renamed=True) 

68 verifier.require_compressed("visit_image", data_id, renamed=True) 

69 # Test difference_image_predetection, because it's downstream of 

70 # future_visit_image in the rewrite pipeline. 

71 verifier.compare_images("difference_image_predetection", data_id, renamed=False) 

72 # Test difference_image vs. future_visit_image. 

73 verifier.compare_images("difference_image", data_id, renamed=True) 

74 verifier.require_compressed("difference_image", data_id, renamed=True) 

75 # dia_source detector is downstream of difference_image; make sure 

76 # it's unchanged as a sanity check. 

77 verifier.compare_tables("dia_source_detector", data_id, ignore={"timeProcessedMjdTai"}) 

78 # Query for and compare datasets with {visit, detector, tract} data 

79 # IDs that are downstream of difference_image and maybe visit_image. 

80 for data_id in verifier.process("object_forced_source_unstandardized", where): 

81 verifier.compare_tables("object_forced_source_unstandardized", data_id) 

82 verifier.compare_tables("dia_object_forced_source_unstandardized", data_id) 

83 

84 

85class RewriteVerifier: 

86 def __init__( 

87 self, 

88 butler: Butler, 

89 *, 

90 old_collection: str, 

91 new_collection: str, 

92 old_prefix: str, 

93 new_prefix: str, 

94 ) -> None: 

95 self.butler = butler 

96 self.old_collection = old_collection 

97 self.new_collection = new_collection 

98 self.old_prefix = old_prefix 

99 self.new_prefix = new_prefix 

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

101 

102 def process( 

103 self, base_name: str, where: str, *dimensions: str, renamed: bool = False 

104 ) -> Iterable[DataCoordinate]: 

105 import tqdm 

106 

107 base_dataset_type = self.butler.get_dataset_type( 

108 f"{self.new_prefix}{base_name}" if renamed else base_name 

109 ) 

110 dimension_group = base_dataset_type.dimensions.union(self.butler.dimensions.conform(dimensions)) 

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

112 data_ids = list( 

113 tqdm.tqdm( 

114 query.where(where) 

115 .join_dataset_search(base_dataset_type.name, collections=[self.new_collection]) 

116 .data_ids(dimension_group), 

117 desc=f"querying for {self.new_prefix}{base_name} in {self.new_collection}.", 

118 ) 

119 ) 

120 random.shuffle(data_ids) 

121 self._progress = tqdm.tqdm(data_ids) 

122 yield from self._progress 

123 

124 def get_old( 

125 self, base_name: str, data_id: DataCoordinate, storage_class: str | None = None, renamed: bool = False 

126 ) -> Any | None: 

127 ref = self.butler.find_dataset( 

128 f"{self.old_prefix}{base_name}" if renamed else base_name, 

129 data_id, 

130 collections=[self.old_collection], 

131 ) 

132 if ref is None: 

133 return None 

134 assert not ref.run.startswith(self.new_collection), f"Bad old run collection for {ref}." 

135 return self.butler.get(ref, storageClass=storage_class) 

136 

137 def get_new( 

138 self, base_name: str, data_id: DataCoordinate, renamed: bool = False, storage_class: str | None = None 

139 ) -> Any | None: 

140 ref = self.butler.find_dataset( 

141 f"{self.new_prefix}{base_name}" if renamed else base_name, 

142 data_id, 

143 collections=[self.new_collection], 

144 ) 

145 if ref is None or ref.run.startswith(self.old_collection): 

146 return None 

147 return self.butler.get(ref, storageClass=storage_class) 

148 

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

150 if self._progress is not None: 

151 self._progress.write(f"{base_name}@{data_id}: {msg}") 

152 else: 

153 print(f"{base_name}@{data_id}: {msg}") 

154 

155 def compare_images(self, base_name: str, data_id: DataCoordinate, renamed: bool = False) -> bool: 

156 old = self.get_old(base_name, data_id, renamed=renamed) 

157 new = self.get_new(base_name, data_id, renamed=renamed, storage_class="ExposureF") 

158 if old is None: 

159 if new is None: 

160 return False 

161 else: 

162 self.print(base_name, data_id, "old does not exist.") 

163 return True 

164 elif new is None: 

165 self.print(base_name, data_id, "new does not exist.") 

166 return True 

167 differences = [] 

168 if not np.array_equal(old.image.array, new.image.array, equal_nan=True): 

169 differences.append("image") 

170 if not np.array_equal(old.mask.array, new.mask.array): 

171 differences.append("mask") 

172 self.print_mask_diff(base_name, data_id, old.mask, new.mask) 

173 if not np.array_equal(old.variance.array, new.variance.array): 

174 differences.append("variance") 

175 if differences: 

176 self.print(base_name, data_id, ", ".join(differences)) 

177 return bool(differences) 

178 

179 def compare_tables(self, base_name: str, data_id: DataCoordinate, ignore: Set[str] = frozenset()) -> bool: 

180 old = self.get_old(base_name, data_id, storage_class="ArrowAstropy") 

181 new = self.get_new(base_name, data_id, storage_class="ArrowAstropy") 

182 if old is None: 

183 if new is None: 

184 return False 

185 else: 

186 self.print(base_name, data_id, "old does not exist.") 

187 return True 

188 elif new is None: 

189 self.print(base_name, data_id, "new does not exist.") 

190 return True 

191 if len(old) != len(new): 

192 self.print(base_name, data_id, f"row counts differ ({len(old) != len(new)})") 

193 return True 

194 if old.colnames != new.colnames: 

195 self.print(base_name, data_id, "column names differ") 

196 return True 

197 differences = [] 

198 for name in old.colnames: 

199 if name in ignore: 

200 continue 

201 try: 

202 if not np.array_equal( 

203 old[name], new[name], equal_nan=issubclass(old[name].dtype.type, np.floating) 

204 ): 

205 differences.append(name) 

206 except Exception as err: 

207 err.add_note(name) 

208 raise 

209 if differences: 

210 self.print(base_name, data_id, ", ".join(differences)) 

211 return bool(differences) 

212 

213 def require_compressed(self, base_name: str, data_id: DataCoordinate, renamed: bool = False) -> None: 

214 ref = self.butler.find_dataset( 

215 f"{self.new_prefix}{base_name}" if renamed else base_name, 

216 data_id, 

217 collections=[self.new_collection], 

218 ) 

219 if ref is None: 

220 return 

221 path = self.butler.getURI(ref) 

222 fs: fsspec.AbstractFileSystem 

223 fs, fp = path.to_fsspec() 

224 with fs.open(fp) as stream: 

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

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

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

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

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

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

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

232 

233 def print_log_diff(self, task_label: str, data_id: DataCoordinate) -> None: 

234 old = self.get_old(f"{task_label}_log", data_id) 

235 new = self.get_new(f"{task_label}_log", data_id) 

236 assert old is not None 

237 assert new is not None 

238 with open(f"{task_label}-{'_'.join(str(v) for v in data_id.required_values)}.diff", "w") as stream: 

239 for line in difflib.unified_diff([r.message for r in old], [r.message for r in new]): 

240 stream.write(line) 

241 stream.write("\n") 

242 

243 def print_mask_diff(self, base_name: str, data_id: DataCoordinate, old_mask: Any, new_mask: Any) -> None: 

244 bitmasks: dict[str, int] = { 

245 name: old_mask.getPlaneBitMask(name) for name in old_mask.getMaskPlaneDict() 

246 } 

247 

248 def interpret_mask_value(v: int) -> list[str]: 

249 return [name for name, bitmask in bitmasks.items() if (bitmask & v)] 

250 

251 with open(f"{base_name}-{'_'.join(str(v) for v in data_id.required_values)}.mask", "w") as stream: 

252 for value, count in zip(*np.unique_counts(new_mask.array & ~old_mask.array)): 

253 if not value: 

254 continue 

255 planes = interpret_mask_value(value) 

256 print(f"+ {{{', '.join(planes)}}}: {count}", file=stream) 

257 for value, count in zip(*np.unique_counts(old_mask.array & ~new_mask.array)): 

258 if not value: 

259 continue 

260 planes = interpret_mask_value(value) 

261 print(f"- {{{', '.join(planes)}}}: {count}", file=stream) 

262 

263 

264if __name__ == "__main__": 

265 verify_rewrite()