Coverage for python/lsst/scarlet/lite/io/blend.py: 97%

1from __future__ import annotations 

2 

3import logging 

4from dataclasses import dataclass 

5from functools import cached_property 

6from typing import Any 

7 

8import numpy as np 

9from deprecated.sphinx import deprecated # type: ignore 

10from numpy.typing import DTypeLike 

11 

12from ..bbox import Box 

13from ..blend import Blend 

14from ..observation import Observation 

15from .blend_base import ScarletBlendBaseData 

16from .migration import PRE_SCHEMA, MigrationRegistry, migration 

17from .source import ScarletSourceBaseData 

18from .utils import decode_metadata, encode_metadata, extract_from_metadata 

19 

20__all__ = ["ScarletBlendData"] 

21 

22CURRENT_SCHEMA = "1.0.0" 

23BLEND_TYPE = "blend" 

24MigrationRegistry.set_current(BLEND_TYPE, CURRENT_SCHEMA) 

25logger = logging.getLogger(__name__) 

26 

27 

28@dataclass(kw_only=True) 

29class ScarletBlendData(ScarletBlendBaseData): 

30 """Data for an entire blend. 

31 

32 Attributes 

33 ---------- 

34 blend_type : 

35 The type of blend being stored. 

36 metadata : 

37 Metadata associated with the blend, 

38 for example the order of bands, the PSF, etc. 

39 origin : 

40 The lower bound of the blend's bounding box. 

41 shape : 

42 The shape of the blend's bounding box. 

43 sources : 

44 Data for the sources contained in the blend, 

45 indexed by the source id. 

46 version : 

47 The schema version of the stored data. 

48 """ 

49 

50 blend_type: str = BLEND_TYPE 

51 origin: tuple[int, int] 

52 shape: tuple[int, int] 

53 sources: dict[Any, ScarletSourceBaseData] 

54 version: str = CURRENT_SCHEMA 

55 

56 @cached_property 

57 def bbox(self) -> Box: 

58 """The bounding box of the blend""" 

59 return Box(self.shape, origin=self.origin) 

60 

61 def as_dict(self) -> dict: 

62 """Return the object encoded into a dict for JSON serialization 

63 

64 Returns 

65 ------- 

66 result : 

67 The object encoded as a JSON compatible dict 

68 """ 

69 result: dict[str, Any] = { 

70 "blend_type": self.blend_type, 

71 "origin": self.origin, 

72 "shape": self.shape, 

73 "sources": {bid: source.as_dict() for bid, source in self.sources.items()}, 

74 "version": self.version, 

75 } 

76 if self.metadata is not None: 

77 result["metadata"] = encode_metadata(self.metadata) 

78 return result 

79 

80 @classmethod 

81 def from_dict(cls, data: dict, dtype: DTypeLike = np.float32) -> ScarletBlendData: 

82 """Reconstruct `ScarletBlendData` from JSON compatible 

83 dict. 

84 

85 Parameters 

86 ---------- 

87 data : 

88 Dictionary representation of the object 

89 dtype : 

90 Datatype of the resulting model. 

91 

92 Returns 

93 ------- 

94 result : 

95 The reconstructed object 

96 """ 

97 data = MigrationRegistry.migrate(BLEND_TYPE, data) 

98 metadata = data.get("metadata", None) 

99 

100 return cls( 

101 origin=tuple(data["origin"]), # type: ignore 

102 shape=tuple(data["shape"]), # type: ignore 

103 sources={ 

104 bid: ScarletSourceBaseData.from_dict(source, dtype=dtype) 

105 for bid, source in data["sources"].items() 

106 }, 

107 metadata=decode_metadata(metadata), 

108 ) 

109 

110 def minimal_data_to_blend( 

111 self, 

112 model_psf: np.ndarray | None = None, 

113 psf: np.ndarray | None = None, 

114 bands: tuple[str] | None = None, 

115 dtype: DTypeLike = np.float32, 

116 ) -> Blend: 

117 """Convert the storage data model into a scarlet lite blend 

118 

119 Parameters 

120 ---------- 

121 model_psf : 

122 PSF in model space (usually a nyquist sampled circular Gaussian). 

123 psf : 

124 The PSF of the observation. 

125 If not provided, the PSF stored in the blend data is used. 

126 bands : 

127 The bands in the blend model. 

128 If not provided, the bands stored in the blend data are used. 

129 dtype : 

130 The data type of the model that is generated. 

131 

132 Returns 

133 ------- 

134 blend : 

135 A scarlet blend model extracted from persisted data. 

136 """ 

137 

138 _model_psf: np.ndarray = extract_from_metadata(model_psf, self.metadata, "model_psf") 

139 _psf: np.ndarray = extract_from_metadata(psf, self.metadata, "psf") 

140 _bands: tuple[str] = extract_from_metadata(bands, self.metadata, "bands") 

141 model_box = self.bbox 

142 observation = Observation.empty( 

143 bands=_bands, 

144 psfs=_psf, 

145 model_psf=_model_psf, 

146 bbox=model_box, 

147 dtype=dtype, 

148 ) 

149 return self.to_blend(observation) 

150 

151 def to_blend(self, observation: Observation) -> Blend: 

152 """Convert the storage data model into a scarlet lite blend 

153 

154 Parameters 

155 ---------- 

156 observation : 

157 The observation that contains the blend. 

158 If `observation` is ``None`` then an `Observation` containing 

159 no image data is initialized. 

160 

161 Returns 

162 ------- 

163 blend : 

164 A scarlet blend model extracted from persisted data. 

165 """ 

166 sources = [] 

167 for sid, source_data in self.sources.items(): 

168 source = source_data.to_source(observation) 

169 # Ensure that the source id is persisted to its metadata 

170 if source.metadata is None: 

171 source.metadata = {} 

172 if "id" not in source.metadata: 

173 source.metadata["id"] = sid 

174 sources.append(source) 

175 

176 return Blend(sources=sources, observation=observation, metadata=self.metadata) 

177 

178 @staticmethod 

179 @deprecated( 

180 reason="ScarletBlendData.from_blend is deprecated. Use blend.to_data() instead.", 

181 version="v30.0", 

182 category=FutureWarning, 

183 ) 

184 def from_blend(blend: Blend) -> ScarletBlendData: 

185 """Deprecated: Convert a scarlet lite blend into a storage data model. 

186 

187 Parameters 

188 ---------- 

189 blend : 

190 The blend to convert. 

191 Returns 

192 ------- 

193 result : 

194 The storage data model representing the blend. 

195 """ 

196 return blend.to_data() 

197 

198 

199ScarletBlendData.register() 

200 

201 

202@migration(BLEND_TYPE, PRE_SCHEMA) 

203def _to_1_0_0(data: dict) -> dict: 

204 """Migrate a pre-schema blend to schema version 1.0.0 

205 

206 Parameters 

207 ---------- 

208 data : 

209 The data to migrate. 

210 

211 Returns 

212 ------- 

213 result : 

214 The migrated data. 

215 """ 

216 # Support legacy models before metadata was used 

217 if "metadata" not in data and "psf" in data: 217 ↛ 224line 217 didn't jump to line 224 because the condition on line 217 was always true

218 data["metadata"] = { 

219 "psf": data["psf"], 

220 "psf_shape": data["psf_shape"], 

221 "bands": tuple(data["bands"]), 

222 "array" "_keys": ["psf"], 

223 } 

224 data["version"] = "1.0.0" 

225 return data