Coverage for python / lsst / images / serialization / _asdf_utils.py: 68%

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 "ArrayReferenceModel", 

16 "ArrayReferenceQuantityModel", 

17 "InlineArray", 

18 "InlineArrayModel", 

19 "InlineArrayQuantity", 

20 "InlineArrayQuantityModel", 

21 "Quantity", 

22 "QuantityModel", 

23 "Time", 

24 "TimeModel", 

25 "Unit", 

26) 

27 

28from typing import Annotated, Any, Literal 

29 

30import astropy.time 

31import astropy.units 

32import numpy as np 

33import pydantic 

34import pydantic_core.core_schema as pcs 

35 

36from ._dtypes import NumberType 

37 

38 

39class _UnitSerialization: 

40 """Pydantic hooks for unit serialization. 

41 

42 This class provides implementations for the `Unit` type alias for 

43 `astropy.unit.Unit` that adds Pydantic serialization and validation. 

44 """ 

45 

46 @classmethod 

47 def __get_pydantic_core_schema__( 

48 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

49 ) -> pcs.CoreSchema: 

50 from_str_schema = pcs.chain_schema( 

51 [ 

52 pcs.str_schema(), 

53 pcs.no_info_plain_validator_function(cls.from_str), 

54 ] 

55 ) 

56 return pcs.json_or_python_schema( 

57 json_schema=from_str_schema, 

58 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.units.UnitBase), from_str_schema]), 

59 serialization=pcs.plain_serializer_function_ser_schema(cls.to_str), 

60 ) 

61 

62 @classmethod 

63 def from_str(cls, value: str) -> astropy.units.UnitBase: 

64 try: 

65 return astropy.units.Unit(value, format="vounit") 

66 except ValueError: 

67 pass 

68 # Some important units (e.g. "dn") are not supported by vounit, so 

69 # fall back to letting astropy to infer the format. 

70 return astropy.units.Unit(value) 

71 

72 @staticmethod 

73 def to_str(unit: astropy.units.UnitBase) -> str: 

74 try: 

75 return unit.to_string("vounit") 

76 except ValueError: 

77 pass 

78 # Some important units (e.g. "dn") are not supported by vounit, so 

79 # fall back to letting astropy use the default format. 

80 return unit.to_string() 

81 

82 

83type Unit = Annotated[ 

84 astropy.units.UnitBase, 

85 _UnitSerialization, 

86 pydantic.WithJsonSchema( 

87 { 

88 "type": "string", 

89 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

90 "id": "http://stsci.edu/schemas/asdf/unit/unit-1.0.0", 

91 "tag": "!unit/unit-1.0.0", 

92 } 

93 ), 

94] 

95 

96 

97class ArrayReferenceModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

98 """Model for a subset of the ASDF 'ndarray' schema, in the case where the 

99 array data is stored elsewhere. 

100 """ 

101 

102 source: str | int = pydantic.Field(description="Location of the underlying binary data.") 

103 shape: list[int] = pydantic.Field(description="Size of the array in each dimension.") 

104 datatype: NumberType = pydantic.Field(description="Data type of the array.") 

105 byteorder: Literal["big"] = pydantic.Field(default="big", description="Byte order for the binary data.") 

106 

107 def with_units(self, unit: astropy.units.UnitBase) -> ArrayReferenceQuantityModel: 

108 """Add units, transforming this model into a Quantity model.""" 

109 return ArrayReferenceQuantityModel.model_construct(value=self, unit=unit) 

110 

111 model_config = pydantic.ConfigDict( 

112 json_schema_extra={ 

113 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

114 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0", 

115 "tag": "!core/ndarray-1.1.0", 

116 } 

117 ) 

118 

119 

120class InlineArrayModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

121 """Model for a subset of the ASDF 'ndarray' schema, in the case where the 

122 array data is stored inline. 

123 """ 

124 

125 data: list[Any] 

126 datatype: NumberType 

127 

128 @property 

129 def shape(self) -> tuple[int, ...]: 

130 """The shape of the array (`tuple` [`int`, ...]).""" 

131 return self._extract_shape(self.data) 

132 

133 def with_units(self, unit: astropy.unit.UnitBase) -> InlineArrayQuantityModel: 

134 """Add units, transforming this model in to a Quantity model.""" 

135 return InlineArrayQuantityModel.model_construct(value=self, unit=unit) 

136 

137 @classmethod 

138 def _extract_shape(cls, data: list[Any], current: tuple[int, ...] = ()) -> tuple[int, ...]: 

139 if not data: 

140 return current + (0,) 

141 if not isinstance(data[0], list): 

142 return current + (len(data),) 

143 return cls._extract_shape(data[0], current + (len(data),)) 

144 

145 model_config = pydantic.ConfigDict( 

146 json_schema_extra={ 

147 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

148 "id": "http://stsci.edu/schemas/asdf/core/ndarray-1.1.0", 

149 "tag": "!core/ndarray-1.1.0", 

150 } 

151 ) 

152 

153 

154class _InlineArraySerialization: 

155 """Pydantic hooks for array serialization. 

156 

157 This class provides implementations for the `Array` type alias for 

158 `numpy.ndarray` that adds Pydantic serialization and validation. 

159 """ 

160 

161 @classmethod 

162 def __get_pydantic_core_schema__( 

163 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

164 ) -> pcs.CoreSchema: 

165 from_model_schema = pcs.chain_schema( 

166 [ 

167 handler(InlineArrayModel), 

168 pcs.no_info_plain_validator_function(cls.from_model), 

169 ] 

170 ) 

171 return pcs.json_or_python_schema( 

172 json_schema=from_model_schema, 

173 python_schema=pcs.union_schema([pcs.is_instance_schema(np.ndarray), from_model_schema]), 

174 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model), 

175 ) 

176 

177 @classmethod 

178 def from_model(cls, model: InlineArrayModel) -> np.ndarray: 

179 return np.array(model.data, dtype=model.datatype.to_numpy()) 

180 

181 @classmethod 

182 def to_model(cls, array: np.ndarray) -> InlineArrayModel: 

183 datatype = NumberType.from_numpy(array.dtype) 

184 return InlineArrayModel(data=array.tolist(), datatype=datatype) 

185 

186 

187type InlineArray = Annotated[np.ndarray, _InlineArraySerialization] 

188 

189 

190class QuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

191 """Model for a subset of the ASDF 'quantity' schema for scalars.""" 

192 

193 value: pydantic.StrictFloat 

194 unit: Unit 

195 

196 model_config = pydantic.ConfigDict( 

197 json_schema_extra={ 

198 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

199 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", 

200 "tag": "!unit/quantity-1.2.0", 

201 } 

202 ) 

203 

204 

205class InlineArrayQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

206 """Model for a subset of the ASDF 'quantity' schema for inline arrays.""" 

207 

208 value: InlineArrayModel 

209 unit: Unit 

210 

211 model_config = pydantic.ConfigDict( 

212 json_schema_extra={ 

213 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

214 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", 

215 "tag": "!unit/quantity-1.2.0", 

216 } 

217 ) 

218 

219 

220class ArrayReferenceQuantityModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

221 """Model for a subset of the ASDF 'quantity' schema for external arrays.""" 

222 

223 value: ArrayReferenceModel 

224 unit: Unit 

225 

226 model_config = pydantic.ConfigDict( 

227 json_schema_extra={ 

228 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

229 "id": "http://stsci.edu/schemas/asdf/unit/quantity-1.2.0", 

230 "tag": "!unit/quantity-1.2.0", 

231 } 

232 ) 

233 

234 

235class _QuantitySerialization: 

236 """Pydantic hooks for scalar quantity serialization.""" 

237 

238 @classmethod 

239 def __get_pydantic_core_schema__( 

240 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

241 ) -> pcs.CoreSchema: 

242 from_model_schema = pcs.chain_schema( 

243 [ 

244 handler(QuantityModel), 

245 pcs.no_info_plain_validator_function(cls.from_model), 

246 ] 

247 ) 

248 return pcs.json_or_python_schema( 

249 json_schema=from_model_schema, 

250 python_schema=pcs.union_schema( 

251 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema] 

252 ), 

253 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model), 

254 ) 

255 

256 @classmethod 

257 def from_model(cls, model: QuantityModel) -> astropy.units.Quantity: 

258 return astropy.units.Quantity(model.value, unit=model.unit) 

259 

260 @classmethod 

261 def to_model(cls, quantity: astropy.units.Quantity) -> QuantityModel: 

262 assert quantity.isscalar 

263 return QuantityModel(value=quantity.to_value(), unit=quantity.unit) 

264 

265 

266type Quantity = Annotated[astropy.units.Quantity, _QuantitySerialization] 

267 

268 

269class _InlineArrayQuantitySerialization: 

270 """Pydantic hooks for inline array quantity serialization.""" 

271 

272 @classmethod 

273 def __get_pydantic_core_schema__( 

274 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

275 ) -> pcs.CoreSchema: 

276 from_model_schema = pcs.chain_schema( 

277 [ 

278 handler(InlineArrayQuantityModel), 

279 pcs.no_info_plain_validator_function(cls.from_model), 

280 ] 

281 ) 

282 return pcs.json_or_python_schema( 

283 json_schema=from_model_schema, 

284 python_schema=pcs.union_schema( 

285 [pcs.is_instance_schema(astropy.units.Quantity), from_model_schema] 

286 ), 

287 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model), 

288 ) 

289 

290 @classmethod 

291 def from_model(cls, model: InlineArrayQuantityModel) -> astropy.units.Quantity: 

292 return astropy.units.Quantity(_InlineArraySerialization.from_model(model.value), unit=model.unit) 

293 

294 @classmethod 

295 def to_model(cls, quantity: astropy.units.Quantity) -> InlineArrayQuantityModel: 

296 assert quantity.isscalar 

297 return InlineArrayQuantityModel( 

298 value=_InlineArraySerialization.to_model(quantity.to_value()), 

299 unit=quantity.unit, 

300 ) 

301 

302 

303type InlineArrayQuantity = Annotated[astropy.units.Quantity, _InlineArrayQuantitySerialization] 

304 

305 

306class TimeModel(pydantic.BaseModel, ser_json_inf_nan="constants"): 

307 """Model for a subset of the ASDF 'time' schema.""" 

308 

309 value: str 

310 scale: Literal["utc", "tai"] 

311 format: Literal["iso"] = "iso" 

312 

313 model_config = pydantic.ConfigDict( 

314 json_schema_extra={ 

315 "$schema": "http://stsci.edu/schemas/yaml-schema/draft-01", 

316 "id": "http://stsci.edu/schemas/asdf/time/time-1.2.0", 

317 "tag": "!time/time-1.2.0", 

318 } 

319 ) 

320 

321 

322class _TimeSerialization: 

323 """Pydantic hooks for time serialization. 

324 

325 This class provides implementations for the `Time` type alias for 

326 `astropy.time.Time` that adds Pydantic serialization and validation. 

327 """ 

328 

329 @classmethod 

330 def __get_pydantic_core_schema__( 

331 cls, source_type: Any, handler: pydantic.GetCoreSchemaHandler 

332 ) -> pcs.CoreSchema: 

333 from_model_schema = pcs.chain_schema( 

334 [ 

335 TimeModel.__pydantic_core_schema__, 

336 pcs.no_info_plain_validator_function(cls.from_model), 

337 ] 

338 ) 

339 return pcs.json_or_python_schema( 

340 json_schema=from_model_schema, 

341 python_schema=pcs.union_schema([pcs.is_instance_schema(astropy.time.Time), from_model_schema]), 

342 serialization=pcs.plain_serializer_function_ser_schema(cls.to_model, info_arg=False), 

343 ) 

344 

345 @classmethod 

346 def from_model(cls, model: TimeModel) -> astropy.time.Time: 

347 return astropy.time.Time(model.value, scale=model.scale, format=model.format) 

348 

349 @classmethod 

350 def to_model(cls, time: astropy.time.Time) -> TimeModel: 

351 if time.scale != "utc" and time.scale != "tai": 

352 time = time.tai 

353 return TimeModel(value=time.to_value("iso"), scale=time.scale, format="iso") 

354 

355 

356type Time = Annotated[astropy.time.Time, _TimeSerialization]