Coverage for python/lsst/images/cli/_fuzz.py: 84%

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. 

11from __future__ import annotations 

12 

13__all__ = ("fuzz_masked_image", "shuffle_blocks") 

14 

15from pathlib import Path 

16from typing import Any 

17 

18import click 

19import numpy as np 

20 

21from .._geom import YX 

22from ..fits import ( 

23 FitsCompressionAlgorithm, 

24 FitsCompressionOptions, 

25 FitsDitherAlgorithm, 

26 FitsQuantizationOptions, 

27) 

28from ..serialization import ArchiveReadError, backend_for_path, read 

29 

30 

31def shuffle_blocks( 

32 image: np.ndarray, 

33 mask: np.ndarray, 

34 variance: np.ndarray, 

35 block_shape: YX[int], 

36 rng: np.random.Generator, 

37) -> None: 

38 """Shuffle image, mask, and variance pixels within each block in place. 

39 

40 A single permutation is drawn for each block and applied to all three 

41 planes, so a pixel's values move together and stay mutually consistent. 

42 Blocks at the edges may be smaller than ``block_shape`` when its dimensions 

43 do not divide the image evenly. 

44 

45 Parameters 

46 ---------- 

47 image : `numpy.ndarray` 

48 The 2-d image plane, modified in place. 

49 mask : `numpy.ndarray` 

50 The mask plane with a leading ``(ny, nx)`` shape and any trailing axes 

51 (for example a per-pixel byte axis), modified in place. 

52 variance : `numpy.ndarray` 

53 The 2-d variance plane, modified in place. 

54 block_shape : `~lsst.images.YX` 

55 The ``(y, x)`` size of a single block. 

56 rng : `numpy.random.Generator` 

57 Random number generator used to permute pixels. 

58 """ 

59 block_y, block_x = block_shape 

60 n_y, n_x = image.shape 

61 for y0 in range(0, n_y, block_y): 

62 y1 = min(y0 + block_y, n_y) 

63 for x0 in range(0, n_x, block_x): 

64 x1 = min(x0 + block_x, n_x) 

65 count = (y1 - y0) * (x1 - x0) 

66 permutation = rng.permutation(count) 

67 for plane in (image, mask, variance): 

68 block = plane[y0:y1, x0:x1] 

69 flat = block.reshape(count, *block.shape[2:]) 

70 block[...] = flat[permutation].reshape(block.shape) 

71 

72 

73def _output_path(in_path: Path, suffix: str) -> Path: 

74 """Insert ``suffix`` before the file extension, preserving the recognised 

75 ``.fits.gz`` compound extension. 

76 """ 

77 name = in_path.name 

78 if name.endswith(".fits.gz"): 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true

79 return in_path.with_name(name[: -len(".fits.gz")] + suffix + ".fits.gz") 

80 return in_path.with_name(in_path.stem + suffix + in_path.suffix) 

81 

82 

83def _block_shape(obj: Any, tile_override: tuple[int, int] | None) -> YX[int]: 

84 """Decide the shuffle block: the override, else the cell size when the 

85 object has a cell grid, else the full image. 

86 """ 

87 if tile_override is not None: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true

88 return YX(int(tile_override[0]), int(tile_override[1])) 

89 cell_shape = getattr(getattr(obj, "grid", None), "cell_shape", None) 

90 if cell_shape is not None: 90 ↛ 91line 90 didn't jump to line 91 because the condition on line 90 was never true

91 return YX(int(cell_shape.y), int(cell_shape.x)) 

92 shape = obj.image.array.shape 

93 return YX(int(shape[0]), int(shape[1])) 

94 

95 

96def _compression_options( 

97 fits_tile: tuple[int, int] | None, quantize_level: float 

98) -> dict[str, FitsCompressionOptions]: 

99 """Build the per-plane compression profile keyed by logical options_name. 

100 

101 ``image`` and ``variance`` (and noise realizations, which reuse the 

102 ``image`` profile) get lossy RICE with subtractive dithering; ``mask`` gets 

103 lossless GZIP. Every other plane falls back to the lossless GZIP default. 

104 """ 

105 lossy = FitsCompressionOptions( 

106 algorithm=FitsCompressionAlgorithm.RICE_1, 

107 tile_shape=fits_tile, 

108 quantization=FitsQuantizationOptions( 

109 dither=FitsDitherAlgorithm.SUBTRACTIVE_DITHER_2, level=quantize_level 

110 ), 

111 ) 

112 gzip = FitsCompressionOptions( 

113 algorithm=FitsCompressionAlgorithm.GZIP_2, tile_shape=fits_tile, quantization=None 

114 ) 

115 return {"image": lossy, "variance": lossy, "mask": gzip} 

116 

117 

118def _verify(in_path: Path, original: dict[str, np.ndarray], check: Any) -> None: 

119 """Confirm the shuffled planes actually changed in the re-read output.""" 

120 for name in ("image", "variance"): 

121 orig = original[name] 

122 finite = np.isfinite(orig) 

123 if not np.any(finite) or np.ptp(orig[finite]) == 0: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 continue # A constant plane cannot change under permutation. 

125 new = getattr(check, name).array 

126 changed = float(np.mean(new[finite] != orig[finite])) 

127 if changed < 0.5: 127 ↛ 128line 127 didn't jump to line 128 because the condition on line 127 was never true

128 raise click.ClickException( 

129 f"{name} plane barely changed ({changed:.1%}) after fuzzing {in_path}." 

130 ) 

131 orig_mask = original["mask"] 

132 if np.ptp(orig_mask) != 0 and np.array_equal(check.mask.array, orig_mask): 132 ↛ 133line 132 didn't jump to line 133 because the condition on line 132 was never true

133 raise click.ClickException(f"Mask plane unchanged after fuzzing {in_path}.") 

134 

135 

136def _fuzz_file( 

137 in_path: Path, 

138 out_path: Path, 

139 *, 

140 seed: int, 

141 tile_override: tuple[int, int] | None, 

142 quantize_level: float, 

143 compression_seed: int, 

144) -> None: 

145 """Read one file, shuffle its proprietary planes, write it, and verify.""" 

146 try: 

147 obj = read(str(in_path)) 

148 except (ValueError, ArchiveReadError) as err: 

149 raise click.ClickException(f"Could not read {in_path}: {err}") from None 

150 

151 block_shape = _block_shape(obj, tile_override) 

152 original = { 

153 "image": obj.image.array.copy(), 

154 "mask": obj.mask.array.copy(), 

155 "variance": obj.variance.array.copy(), 

156 } 

157 shuffle_blocks( 

158 obj.image.array, obj.mask.array, obj.variance.array, block_shape, np.random.default_rng(seed) 

159 ) 

160 

161 kwargs: dict[str, Any] = {} 

162 if backend_for_path(str(out_path)).name == "fits": 162 ↛ 170line 162 didn't jump to line 170 because the condition on line 162 was always true

163 # tile_shape stays None unless overridden, so each object's natural 

164 # tiling applies (cell for CellCoadd, astropy default otherwise); an 

165 # explicit override aligns the FITS tiles with the shuffle blocks. 

166 fits_tile = block_shape if tile_override is not None else None 

167 kwargs["compression_options"] = _compression_options(fits_tile, quantize_level) 

168 kwargs["compression_seed"] = compression_seed 

169 

170 obj.write(str(out_path), **kwargs) 

171 _verify(in_path, original, read(str(out_path))) 

172 click.echo(f"Fuzzed {in_path} -> {out_path}") 

173 

174 

175@click.command(name="fuzz-masked-image") 

176@click.argument("files", nargs=-1, type=click.Path(exists=True, dir_okay=False, path_type=Path)) 

177@click.option("--seed", type=int, default=1, show_default=True, help="Seed for pixel shuffling.") 

178@click.option( 

179 "--suffix", 

180 default=".fuzzed", 

181 show_default=True, 

182 help="Text inserted before the extension of each output file.", 

183) 

184@click.option( 

185 "--tile-shape", 

186 type=int, 

187 nargs=2, 

188 default=None, 

189 metavar="Y X", 

190 help="Override the shuffle block and FITS compression tile; default is the " 

191 "cell size if the object has a cell grid, else the full image.", 

192) 

193@click.option( 

194 "--quantize-level", 

195 type=float, 

196 default=16.0, 

197 show_default=True, 

198 help="RICE quantization level for the lossy image/variance planes.", 

199) 

200@click.option( 

201 "--compression-seed", 

202 type=int, 

203 default=1, 

204 show_default=True, 

205 help="FITS tile-compression dither seed (FITS output only).", 

206) 

207@click.option("--overwrite", is_flag=True, default=False, help="Overwrite existing output files.") 

208def fuzz_masked_image( 

209 files: tuple[Path, ...], 

210 seed: int, 

211 suffix: str, 

212 tile_shape: tuple[int, int] | None, 

213 quantize_level: float, 

214 compression_seed: int, 

215 overwrite: bool, 

216) -> None: 

217 """Shuffle the proprietary pixels of MaskedImage files for public release. 

218 

219 Each FILE is read in whatever format it is given, its image, mask, and 

220 variance planes are shuffled within tiles (one shared permutation per tile 

221 keeps the three planes mutually consistent), and the result is written 

222 beside the input with SUFFIX inserted before the extension. The output 

223 format follows that extension. Every other plane and all metadata are 

224 written back unchanged, and each output is re-read to confirm the 

225 proprietary planes really changed. 

226 """ 

227 if not files: 

228 raise click.UsageError("No input files given.") 

229 failures = 0 

230 for in_path in files: 

231 out_path = _output_path(in_path, suffix) 

232 if out_path.exists() and not overwrite: 

233 click.echo(f"Skipping {in_path}: output {out_path} already exists.", err=True) 

234 continue 

235 try: 

236 _fuzz_file( 

237 in_path, 

238 out_path, 

239 seed=seed, 

240 tile_override=tile_shape if tile_shape else None, 

241 quantize_level=quantize_level, 

242 compression_seed=compression_seed, 

243 ) 

244 except click.ClickException as err: 

245 click.echo(f"Error: {err.format_message()}", err=True) 

246 failures += 1 

247 if failures: 247 ↛ 248line 247 didn't jump to line 248 because the condition on line 247 was never true

248 raise click.ClickException(f"{failures} file(s) failed to fuzz.")