Coverage for python / lsst / images / fits / formatters.py: 47%

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__ = ("GenericFormatter", "ImageFormatter", "MaskedImageFormatter", "VisitImageFormatter") 

15 

16import enum 

17import hashlib 

18import json 

19from typing import Any, ClassVar 

20 

21import astropy.io.fits 

22 

23from lsst.daf.butler import DatasetProvenance, FormatterV2 

24from lsst.resources import ResourcePath 

25 

26from .._geom import Box 

27from .._image import Image 

28from .._mask import Mask 

29from .._masked_image import MaskedImageSerializationModel 

30from .._transforms import Projection, ProjectionSerializationModel 

31from .._visit_image import VisitImageSerializationModel 

32from ..serialization import TableCellReferenceModel 

33from ._common import FitsCompressionOptions 

34from ._input_archive import FitsInputArchive, read 

35from ._output_archive import write 

36 

37 

38class GenericFormatter(FormatterV2): 

39 """The butler interface to FITS archive serialization. 

40 

41 Serialized types must meet all the requirements of the `read` and `write` 

42 functions. 

43 

44 Notes 

45 ----- 

46 This formatter just forwards all read parameters it receives as 

47 ``**kwargs`` to `.read` and hence the ``deserialize`` method of the type it 

48 is reading. This may or may not be appropriate. 

49 

50 This formatter must be subclassed to add component support. 

51 

52 The write parameter configuration for this formatter is designed to be 

53 identical to that for the legacy FITS formatters defined in 

54 `lsst.obs.base`. 

55 """ 

56 

57 default_extension: ClassVar[str] = ".fits" 

58 can_read_from_uri: ClassVar[bool] = True 

59 supported_write_parameters: ClassVar[frozenset[str]] = frozenset({"recipe"}) 

60 

61 butler_provenance: DatasetProvenance | None = None 

62 

63 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any: 

64 pytype = self.dataset_ref.datasetType.storageClass.pytype 

65 kwargs = self.file_descriptor.parameters or {} 

66 return read(pytype, uri, **kwargs) 

67 

68 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None: 

69 write( 

70 in_memory_dataset, 

71 uri.ospath, 

72 update_header=self._update_header, 

73 compression_options=self._get_compression_options(), 

74 compression_seed=self._get_compression_seed(), 

75 ) 

76 

77 def add_provenance( 

78 self, in_memory_dataset: Any, /, *, provenance: DatasetProvenance | None = None 

79 ) -> Any: 

80 # Instead of attaching the provenance to the object we remember it on 

81 # the formatter, since a Formatter instance is only used once. 

82 self.butler_provenance = provenance 

83 return in_memory_dataset 

84 

85 def _get_compression_seed(self) -> int: 

86 # Set the seed based on data ID (all logic here duplicated from 

87 # obs_base). We can't just use 'hash', since like 'set' that's not 

88 # deterministic. And we can't rely on a DimensionPacker because those 

89 # are only defined for certain combinations of dimensions. Doing an MD5 

90 # of the JSON feels like overkill but I don't really see anything much 

91 # simpler. 

92 hash_bytes = hashlib.md5( 

93 json.dumps(list(self.data_id.required_values)).encode(), 

94 usedforsecurity=False, 

95 ).digest() 

96 # And it *really* feels like overkill when we squash that into the [1, 

97 # 10000] range allowed by FITS. 

98 return 1 + int.from_bytes(hash_bytes) % 9999 

99 

100 def _get_compression_options(self) -> dict[str, FitsCompressionOptions]: 

101 recipe = self.write_parameters.get("recipe", "default") 

102 try: 

103 config = self.write_recipes[recipe] 

104 except KeyError: 

105 if recipe == "default": 105 ↛ 108line 105 didn't jump to line 108 because the condition on line 105 was always true

106 # If there's no default recipe just use the software defaults. 

107 return {} 

108 raise RuntimeError(f"Invalid recipe for ImageFormatter: {recipe!r}.") from None 

109 return {k: FitsCompressionOptions.model_validate(v) for k, v in config.items()} 

110 

111 def _update_header(self, header: astropy.io.fits.Header) -> None: 

112 # Logic here largely lifted from lsst.obs.base.utils, which we 

113 # can't use directly for dependency and maybe mapping-type 

114 # (PropertyList vs. astropy) reasons. We assume we can always add 

115 # long cards (astropy will CONTINUE them) but not comments 

116 # (astropy will truncate and warn on long cards). 

117 for key in list(header): 

118 if key.startswith("LSST BUTLER"): 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true

119 del header[key] 

120 if self.butler_provenance is not None: 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true

121 for key, value in self.butler_provenance.to_flat_dict( 

122 self.dataset_ref, prefix="LSST BUTLER", sep=" ", simple_types=True, max_inputs=3_000 

123 ).items(): 

124 header.set(key, value) 

125 

126 

127class ComponentSentinel(enum.Enum): 

128 """Special values returned by `ImageFormatter.read_component`.""" 

129 

130 UNRECOGNIZED_COMPONENT = enum.auto() 

131 """This formatter does not recognize the given component, but a subclass 

132 might. 

133 """ 

134 

135 INVALID_COMPONENT_MODEL = enum.auto() 

136 """This formatter recognizes the given component, but the expected 

137 attribute of the top-level `..serialization.ArchiveTree` did not exist 

138 or had the wrong type. 

139 """ 

140 

141 

142class ImageFormatter(GenericFormatter): 

143 """The specialized butler interface to FITS archive serialization of 

144 image-like objects with ``projection`` and ``bbox`` components. 

145 

146 Notes 

147 ----- 

148 This formatter works by assuming the `..serialization.ArchiveTree` for the 

149 top-level object has a ``projection`` attribute (a 

150 `..ProjectionSerializationModel`) and a ``bbox`` property (a `..Box`). 

151 

152 Subclasses can add support for additional components by overriding 

153 `read_component`, delegating to `super`, and handling the cases where it 

154 returns a `ComponentSentinel` instance. 

155 """ 

156 

157 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any: 

158 pytype: Any = self.file_descriptor.storageClass.pytype 

159 if component is None: 159 ↛ 162line 159 didn't jump to line 162 because the condition on line 159 was always true

160 result = read(pytype, uri, bbox=self.pop_bbox_from_parameters()) 

161 else: 

162 with FitsInputArchive.open(uri, partial=True) as archive: 

163 tree = archive.get_tree(pytype._get_archive_tree_type(TableCellReferenceModel)) 

164 result = self.read_component(component, tree, archive) 

165 if result is ComponentSentinel.UNRECOGNIZED_COMPONENT: 

166 raise NotImplementedError( 

167 f"Unrecognized component {component!r} for {type(self).__name__}." 

168 ) 

169 if result is ComponentSentinel.INVALID_COMPONENT_MODEL: 

170 raise NotImplementedError( 

171 f"Invalid serialization model for component {component!r} for {type(self).__name__}." 

172 ) 

173 self.check_unhandled_parameters() 

174 return result 

175 

176 def pop_bbox_from_parameters(self) -> Box | None: 

177 parameters = self.file_descriptor.parameters or {} 

178 return parameters.pop("bbox", None) 

179 

180 def check_unhandled_parameters(self) -> None: 

181 if self.file_descriptor.parameters: 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true

182 raise RuntimeError(f"Parameters {list(self.file_descriptor.parameters.keys())} not recognized.") 

183 

184 def read_component( 

185 self, 

186 component: str, 

187 tree: Any, 

188 archive: FitsInputArchive, 

189 ) -> Any: 

190 match component: 

191 case "projection": 

192 if isinstance( 

193 serialized_projection := getattr(tree, "projection", None), ProjectionSerializationModel 

194 ): 

195 return Projection.deserialize(serialized_projection, archive) 

196 else: 

197 return ComponentSentinel.INVALID_COMPONENT_MODEL 

198 case "bbox": 

199 if isinstance(bbox := getattr(tree, "bbox", None), Box): 

200 return bbox 

201 else: 

202 return ComponentSentinel.INVALID_COMPONENT_MODEL 

203 return ComponentSentinel.UNRECOGNIZED_COMPONENT 

204 

205 

206class MaskedImageFormatter(ImageFormatter): 

207 """A specialized butler interface to FITS archive serialization of 

208 the `..MaskedImage` class. 

209 """ 

210 

211 def read_component( 

212 self, 

213 component: str, 

214 tree: Any, 

215 archive: FitsInputArchive, 

216 ) -> Any: 

217 match super().read_component(component, tree, archive): 

218 case ComponentSentinel(): 

219 pass 

220 case handled: 

221 return handled 

222 if not isinstance(tree, MaskedImageSerializationModel): 

223 return ComponentSentinel.INVALID_COMPONENT_MODEL 

224 match component: 

225 case "image": 

226 return Image.deserialize(tree.image, archive, bbox=self.pop_bbox_from_parameters()) 

227 case "mask": 

228 return Mask.deserialize(tree.mask, archive, bbox=self.pop_bbox_from_parameters()) 

229 case "variance": 

230 return Image.deserialize(tree.variance, archive, bbox=self.pop_bbox_from_parameters()) 

231 return ComponentSentinel.UNRECOGNIZED_COMPONENT 

232 

233 

234class VisitImageFormatter(MaskedImageFormatter): 

235 """A specialized butler interface to FITS archive serialization of 

236 the `..VisitImage` class. 

237 """ 

238 

239 def read_component( 

240 self, 

241 component: str, 

242 tree: Any, 

243 archive: FitsInputArchive, 

244 ) -> Any: 

245 match super().read_component(component, tree, archive): 

246 case ComponentSentinel(): 

247 pass 

248 case handled: 

249 return handled 

250 if not isinstance(tree, VisitImageSerializationModel): 

251 return ComponentSentinel.INVALID_COMPONENT_MODEL 

252 match component: 

253 case "psf": 

254 return tree.deserialize_psf(archive) 

255 return ComponentSentinel.UNRECOGNIZED_COMPONENT