Coverage for python/lsst/scarlet/lite/io/cube_component.py: 62%

1from __future__ import annotations 

2 

3from dataclasses import dataclass 

4from typing import Callable 

5 

6import numpy as np 

7from numpy.typing import DTypeLike 

8 

9from ..bbox import Box 

10from ..component import Component 

11from ..image import Image 

12from ..observation import Observation 

13from .component import ScarletComponentBaseData 

14from .migration import PRE_SCHEMA, MigrationRegistry, migration 

15 

16__all__ = ["ScarletCubeComponentData", "ComponentCube"] 

17 

18CURRENT_SCHEMA = "1.0.0" 

19COMPONENT_TYPE = "cube" 

20MigrationRegistry.set_current(COMPONENT_TYPE, CURRENT_SCHEMA) 

21 

22 

23class ComponentCube(Component): 

24 """Dummy component for a component cube. 

25 

26 This is duck-typed to a `lsst.scarlet.lite.Component` in order to 

27 generate a model from the component but it is currently not functional 

28 in that it cannot be optimized, only persisted and loaded. 

29 

30 If scarlet lite ever implements a component as a data cube, 

31 this class can be removed. 

32 """ 

33 

34 def __init__(self, model: Image, peak: tuple[int, int]): 

35 """Initialization 

36 

37 Parameters 

38 ---------- 

39 bands : 

40 model : 

41 The 3D (bands, y, x) model of the component. 

42 peak : 

43 The `(y, x)` peak of the component. 

44 bbox : 

45 The bounding box of the component. 

46 """ 

47 super().__init__(model.bands, model.bbox) 

48 self._model = model 

49 self.peak = peak 

50 

51 def get_model(self) -> Image: 

52 """Generate the model for the source 

53 

54 Returns 

55 ------- 

56 model : 

57 The model as a 3D `(band, y, x)` array. 

58 """ 

59 return self._model 

60 

61 def resize(self, model_box: Box) -> bool: 

62 """Test whether or not the component needs to be resized""" 

63 return False 

64 

65 def update(self, it: int, input_grad: np.ndarray) -> None: 

66 """Implementation of unused abstract method""" 

67 

68 def parameterize(self, parameterization: Callable) -> None: 

69 """Implementation of unused abstract method""" 

70 

71 def to_data(self) -> ScarletCubeComponentData: 

72 """Convert the component to persistable ScarletComponentData 

73 

74 Returns 

75 ------- 

76 component_data: ScarletComponentData 

77 The data object containing the component information 

78 """ 

79 return ScarletCubeComponentData( 

80 origin=self.bbox.origin, # type: ignore 

81 peak=self.peak, # type: ignore 

82 model=self.get_model().data, 

83 ) 

84 

85 

86@dataclass(kw_only=True) 

87class ScarletCubeComponentData(ScarletComponentBaseData): 

88 """Data for a component expressed as a 3D data cube 

89 

90 This is used for scarlet component models that are not factorized, 

91 storing their entire model as a 3D data cube (bands, y, x). 

92 

93 Attributes 

94 ---------- 

95 origin : 

96 The lower bound of the components bounding box. 

97 peak : 

98 The peak of the component. 

99 model : 

100 The model for the component. 

101 """ 

102 

103 origin: tuple[int, int] 

104 peak: tuple[float, float] 

105 model: np.ndarray 

106 component_type: str = COMPONENT_TYPE 

107 version: str = CURRENT_SCHEMA 

108 

109 @property 

110 def shape(self): 

111 return self.model.shape[-2:] 

112 

113 def to_component(self, observation: Observation) -> ComponentCube: 

114 """Convert the storage data model into a scarlet Component 

115 

116 Parameters 

117 ---------- 

118 observation : 

119 The observation that the component is associated with 

120 

121 Returns 

122 ------- 

123 component : 

124 A scarlet component extracted from persisted data. 

125 """ 

126 bbox = Box(self.shape, origin=self.origin) 

127 model = self.model 

128 if self.peak is None: 

129 peak = None 

130 else: 

131 peak = (int(np.round(self.peak[0])), int(np.round(self.peak[0]))) 

132 assert peak is not None 

133 component = ComponentCube( 

134 model=Image(model, yx0=bbox.origin, bands=observation.bands), # type: ignore 

135 peak=peak, 

136 ) 

137 return component 

138 

139 def as_dict(self) -> dict: 

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

141 

142 Returns 

143 ------- 

144 result : 

145 The object encoded as a JSON compatible dict 

146 """ 

147 return { 

148 "origin": self.origin, 

149 "shape": self.model.shape, 

150 "peak": self.peak, 

151 "model": tuple(self.model.flatten().astype(float)), 

152 "component_type": "component", 

153 "version": self.version, 

154 } 

155 

156 @classmethod 

157 def from_dict(cls, data: dict, dtype: DTypeLike | None = None) -> ScarletCubeComponentData: 

158 """Reconstruct `ScarletComponentData` from JSON compatible dict 

159 

160 Parameters 

161 ---------- 

162 data : 

163 Dictionary representation of the object 

164 dtype : 

165 Datatype of the resulting model. 

166 

167 Returns 

168 ------- 

169 result : 

170 The reconstructed object 

171 """ 

172 data = MigrationRegistry.migrate(COMPONENT_TYPE, data) 

173 shape = tuple(data["shape"]) 

174 return cls( 

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

176 peak=data["peak"], 

177 model=np.array(data["model"]).reshape(shape).astype(dtype), 

178 ) 

179 

180 

181ScarletCubeComponentData.register() 

182 

183 

184@migration(COMPONENT_TYPE, PRE_SCHEMA) 

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

186 """Migrate a pre-schema CubeComponent to schema version 1.0.0 

187 

188 There were no changes to this data model in v1.0.0 but we need 

189 to provide a way to migrate pre-schema data. 

190 

191 Parameters 

192 ---------- 

193 data : 

194 The data to migrate. 

195 

196 Returns 

197 ------- 

198 result : 

199 The migrated data. 

200 """ 

201 data["version"] = "1.0.0" 

202 return data