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

1from __future__ import annotations 

2 

3import logging 

4from dataclasses import dataclass 

5from typing import Any 

6 

7import numpy as np 

8from numpy.typing import DTypeLike 

9 

10from ..bbox import Box 

11from ..blend import Blend 

12from ..observation import Observation 

13from .blend_base import ScarletBlendBaseData 

14from .migration import PRE_SCHEMA, MigrationRegistry, migration 

15from .source import ScarletSourceBaseData 

16from .utils import decode_metadata, encode_metadata, extract_from_metadata 

17 

18__all__ = ["ScarletBlendData"] 

19 

20CURRENT_SCHEMA = "1.0.0" 

21BLEND_TYPE = "blend" 

22MigrationRegistry.set_current(BLEND_TYPE, CURRENT_SCHEMA) 

23logger = logging.getLogger(__name__) 

24 

25 

26@dataclass(kw_only=True) 

27class ScarletBlendData(ScarletBlendBaseData): 

28 """Data for an entire blend. 

29 

30 Attributes 

31 ---------- 

32 blend_type : 

33 The type of blend being stored. 

34 metadata : 

35 Metadata associated with the blend, 

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

37 origin : 

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

39 shape : 

40 The shape of the blend's bounding box. 

41 sources : 

42 Data for the sources contained in the blend, 

43 indexed by the source id. 

44 version : 

45 The schema version of the stored data. 

46 """ 

47 

48 blend_type: str = BLEND_TYPE 

49 origin: tuple[int, int] 

50 shape: tuple[int, int] 

51 sources: dict[Any, ScarletSourceBaseData] 

52 version: str = CURRENT_SCHEMA 

53 

54 def as_dict(self) -> dict: 

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

56 

57 Returns 

58 ------- 

59 result : 

60 The object encoded as a JSON compatible dict 

61 """ 

62 result: dict[str, Any] = { 

63 "blend_type": self.blend_type, 

64 "origin": self.origin, 

65 "shape": self.shape, 

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

67 "version": self.version, 

68 } 

69 if self.metadata is not None: 

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

71 return result 

72 

73 @classmethod 

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

75 """Reconstruct `ScarletBlendData` from JSON compatible 

76 dict. 

77 

78 Parameters 

79 ---------- 

80 data : 

81 Dictionary representation of the object 

82 dtype : 

83 Datatype of the resulting model. 

84 

85 Returns 

86 ------- 

87 result : 

88 The reconstructed object 

89 """ 

90 data = MigrationRegistry.migrate(BLEND_TYPE, data) 

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

92 

93 return cls( 

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

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

96 sources={ 

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

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

99 }, 

100 metadata=decode_metadata(metadata), 

101 ) 

102 

103 def minimal_data_to_blend( 

104 self, 

105 model_psf: np.ndarray | None = None, 

106 psf: np.ndarray | None = None, 

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

108 dtype: DTypeLike = np.float32, 

109 ) -> Blend: 

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

111 

112 Parameters 

113 ---------- 

114 model_psf : 

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

116 psf : 

117 The PSF of the observation. 

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

119 bands : 

120 The bands in the blend model. 

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

122 dtype : 

123 The data type of the model that is generated. 

124 

125 Returns 

126 ------- 

127 blend : 

128 A scarlet blend model extracted from persisted data. 

129 """ 

130 

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

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

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

134 model_box = Box(self.shape, origin=self.origin) 

135 observation = Observation.empty( 

136 bands=_bands, 

137 psfs=_psf, 

138 model_psf=_model_psf, 

139 bbox=model_box, 

140 dtype=dtype, 

141 ) 

142 return self.to_blend(observation) 

143 

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

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

146 

147 Parameters 

148 ---------- 

149 observation : 

150 The observation that contains the blend. 

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

152 no image data is initialized. 

153 

154 Returns 

155 ------- 

156 blend : 

157 A scarlet blend model extracted from persisted data. 

158 """ 

159 sources = [] 

160 for source_data in self.sources.values(): 

161 source = source_data.to_source(observation) 

162 sources.append(source) 

163 

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

165 

166 @staticmethod 

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

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

169 

170 Parameters 

171 ---------- 

172 blend : 

173 The blend to convert. 

174 Returns 

175 ------- 

176 result : 

177 The storage data model representing the blend. 

178 """ 

179 logger.warning("ScarletBlendData.from_blend is deprecated. Use blend.to_data() instead.") 

180 return blend.to_data() 

181 

182 

183ScarletBlendData.register() 

184 

185 

186@migration(BLEND_TYPE, PRE_SCHEMA) 

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

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

189 

190 Parameters 

191 ---------- 

192 data : 

193 The data to migrate. 

194 

195 Returns 

196 ------- 

197 result : 

198 The migrated data. 

199 """ 

200 # Support legacy models before metadata was used 

201 if "metadata" not in data and "psf" in data: 

202 data["metadata"] = { 

203 "psf": data["psf"], 

204 "psf_shape": data["psf_shape"], 

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

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

207 } 

208 data["version"] = "1.0.0" 

209 return data