Coverage for python/lsst/images/cli/_convert.py: 18%

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__ = ("convert",) 

14 

15import os 

16import tempfile 

17from typing import TYPE_CHECKING, Any 

18 

19import astropy.io.fits 

20import click 

21from click.core import ParameterSource 

22 

23from ..serialization import backend_for_path, write 

24 

25if TYPE_CHECKING: 

26 from .. import VisitImage 

27 from ..cells import CellCoadd 

28 

29_LEGACY_TYPES = ("visit_image", "cell_coadd") 

30 

31 

32def detect_legacy_type(path: str) -> str | None: 

33 """Return ``"visit_image"`` or ``"cell_coadd"`` from a legacy FITS file's 

34 ``HIERARCH LSST BUTLER DATASETTYPE`` header, or `None` if it cannot be 

35 determined. 

36 

37 A dataset type ending in ``visit_image`` (e.g. ``visit_image``, 

38 ``preliminary_visit_image``, difference images) is a `VisitImage`; one 

39 containing ``coadd`` is a `CellCoadd`. astropy exposes the 

40 ``HIERARCH LSST BUTLER DATASETTYPE`` card as 

41 ``header["LSST BUTLER DATASETTYPE"]``. 

42 """ 

43 dataset_type: str | None = None 

44 with astropy.io.fits.open(path) as hdul: 

45 for hdu in hdul: 

46 value = hdu.header.get("LSST BUTLER DATASETTYPE") 

47 if value: 

48 dataset_type = str(value) 

49 break 

50 if dataset_type is None: 

51 return None 

52 if dataset_type.endswith("visit_image"): 

53 return "visit_image" 

54 if dataset_type == "difference_image": 

55 return "difference_image" 

56 if "coadd" in dataset_type: 

57 return "cell_coadd" 

58 return None 

59 

60 

61def _load_skymap(skymap: str | None, butler: str | None, collection: str | None, skymap_name: str) -> Any: 

62 """Load a skymap object from a pickle path or a butler repository.""" 

63 if skymap is not None: 

64 import pickle 

65 

66 with open(skymap, "rb") as stream: 

67 return pickle.load(stream) 

68 if butler is not None: 

69 if collection is None: 

70 raise click.ClickException("--butler also requires --collection (the skymap's collection).") 

71 from lsst.daf.butler import Butler 

72 

73 with Butler.from_config(butler) as repo: 

74 return repo.get("skyMap", skymap=skymap_name, collections=collection) 

75 raise click.ClickException("Converting a cell coadd requires --skymap (a pickled skymap) or --butler.") 

76 

77 

78def _read_legacy( 

79 input: str, 

80 legacy_type: str, 

81 skymap: str | None, 

82 butler: str | None, 

83 collection: str | None, 

84 preserve_quantization: bool = False, 

85) -> VisitImage | CellCoadd: 

86 """Read a legacy FITS file into the corresponding lsst.images object.""" 

87 if legacy_type == "visit_image": 

88 from .. import VisitImage 

89 

90 return VisitImage.read_legacy(input, preserve_quantization=preserve_quantization) 

91 if legacy_type == "difference_image": 

92 from .. import DifferenceImage 

93 

94 return DifferenceImage.read_legacy(input, preserve_quantization=preserve_quantization) 

95 if legacy_type == "cell_coadd": 

96 from lsst.cell_coadds import MultipleCellCoadd 

97 

98 from .. import get_legacy_deep_coadd_mask_planes 

99 from ..cells import CellCoadd 

100 

101 legacy = MultipleCellCoadd.read_fits(input) 

102 sky = _load_skymap(skymap, butler, collection, legacy.identifiers.skymap) 

103 tract_info = sky[legacy.identifiers.tract] 

104 return CellCoadd.from_legacy( 

105 legacy, 

106 tract_info=tract_info, 

107 plane_map=get_legacy_deep_coadd_mask_planes(), 

108 ) 

109 raise click.ClickException(f"Conversion of {legacy_type!r} is not yet implemented.") 

110 

111 

112@click.command(name="convert") 

113@click.argument("input", type=click.Path(exists=True, dir_okay=False)) 

114@click.argument("output", type=click.Path(dir_okay=False)) 

115@click.option( 

116 "--type", 

117 "type_", 

118 type=click.Choice(_LEGACY_TYPES), 

119 default=None, 

120 help="Legacy input type; overrides auto-detection.", 

121) 

122@click.option( 

123 "--skymap", 

124 type=click.Path(exists=True, dir_okay=False), 

125 default=None, 

126 help="Pickled skymap (required for cell coadds unless --butler is given).", 

127) 

128@click.option( 

129 "--butler", 

130 default=None, 

131 help="Butler repository to resolve the skymap (cell coadds only).", 

132) 

133@click.option( 

134 "--collection", 

135 default=None, 

136 help="Butler collection holding the skymap (required with --butler).", 

137) 

138@click.option("--overwrite", is_flag=True, default=False, help="Overwrite OUTPUT if it exists.") 

139@click.option( 

140 "--preserve-quantization/--no-preserve-quantization", 

141 default=True, 

142 help=( 

143 "Preserve quantization-compressed pixel values so they can be written " 

144 "back out losslessly (visit images only). On by default." 

145 ), 

146) 

147def convert( 

148 input: str, 

149 output: str, 

150 type_: str | None, 

151 skymap: str | None, 

152 butler: str | None, 

153 collection: str | None, 

154 overwrite: bool, 

155 preserve_quantization: bool, 

156) -> None: 

157 """Convert a legacy FITS file to a new lsst.images format. 

158 

159 The output format is chosen from OUTPUT's extension. 

160 """ 

161 try: 

162 backend = backend_for_path(output) 

163 except ValueError as err: 

164 raise click.ClickException(str(err)) from None 

165 

166 legacy_type = type_ or detect_legacy_type(input) 

167 if legacy_type is None: 

168 raise click.ClickException(f"Could not determine the legacy type of {input!r}; pass --type.") 

169 

170 # The flag is on by default, so only object when the user explicitly set 

171 # it for a type that does not support it (only visit images do). 

172 if legacy_type != "visit_image": 

173 source = click.get_current_context().get_parameter_source("preserve_quantization") 

174 if source == ParameterSource.COMMANDLINE: 

175 raise click.ClickException( 

176 f"--preserve-quantization is only valid for visit images, not {legacy_type!r}." 

177 ) 

178 preserve_quantization = False 

179 

180 output_abs = os.path.realpath(output) 

181 if os.path.realpath(input) == output_abs: 

182 raise click.ClickException("INPUT and OUTPUT must be different paths.") 

183 

184 if os.path.exists(output_abs) and not overwrite: 

185 raise click.ClickException(f"{output!r} already exists; pass --overwrite to replace it.") 

186 

187 try: 

188 obj = _read_legacy(input, legacy_type, skymap, butler, collection, preserve_quantization) 

189 except click.ClickException: 

190 raise 

191 except ImportError as err: 

192 raise click.ClickException( 

193 f"Reading a legacy {legacy_type} requires Rubin packages that are not installed: {err}" 

194 ) from None 

195 

196 # Write to a temporary file in the output's directory and move it into 

197 # place only after a successful write, so a read or write failure never 

198 # destroys an existing OUTPUT (the backends refuse to overwrite, so the 

199 # temporary path must not already exist). 

200 output_dir = os.path.dirname(output_abs) 

201 with tempfile.TemporaryDirectory(dir=output_dir) as staging: 

202 staged = os.path.join(staging, os.path.basename(output_abs)) 

203 write(obj, staged) 

204 os.replace(staged, output_abs) 

205 click.echo(f"Wrote {output} ({backend.name}, {legacy_type}).")