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

1from __future__ import annotations 

2 

3import logging 

4from dataclasses import dataclass 

5 

6import numpy as np 

7from deprecated.sphinx import deprecated # type: ignore 

8from numpy.typing import DTypeLike 

9 

10from ..bbox import Box 

11from ..component import CubeComponent 

12from ..image import Image 

13from ..observation import Observation 

14from .component import ScarletComponentBaseData 

15from .migration import PRE_SCHEMA, MigrationRegistry, migration 

16 

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

18 

19CURRENT_SCHEMA = "1.0.0" 

20COMPONENT_TYPE = "cube" 

21MigrationRegistry.set_current(COMPONENT_TYPE, CURRENT_SCHEMA) 

22 

23logger = logging.getLogger(__name__) 

24 

25 

26@deprecated( 

27 reason="ComponentCube is deprecated and will be removed after scarlet_lite v30.0. " 

28 "Please use CubeComponent instead.", 

29 version="scarlet_lite v30.0", 

30 category=FutureWarning, 

31) 

32class ComponentCube(CubeComponent): 

33 """Deprecated, use CubeComponent instead. 

34 

35 Parameters 

36 ---------- 

37 model: 

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

39 peak: 

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

41 """ 

42 

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

44 super().__init__(model=model, peak=peak) 

45 

46 

47@dataclass(kw_only=True) 

48class ScarletCubeComponentData(ScarletComponentBaseData): 

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

50 

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

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

53 

54 Attributes 

55 ---------- 

56 origin : 

57 The lower bound of the components bounding box. 

58 peak : 

59 The peak of the component. 

60 model : 

61 The model for the component. 

62 """ 

63 

64 origin: tuple[int, int] 

65 peak: tuple[float, float] 

66 model: np.ndarray 

67 component_type: str = COMPONENT_TYPE 

68 version: str = CURRENT_SCHEMA 

69 

70 @property 

71 def shape(self): 

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

73 

74 def to_component(self, observation: Observation) -> CubeComponent: 

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

76 

77 Parameters 

78 ---------- 

79 observation : 

80 The observation that the component is associated with 

81 

82 Returns 

83 ------- 

84 component : 

85 A scarlet component extracted from persisted data. 

86 """ 

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

88 model = self.model 

89 if self.peak is None: 

90 peak = None 

91 else: 

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

93 assert peak is not None 

94 component = CubeComponent( 

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

96 peak=peak, 

97 ) 

98 return component 

99 

100 def as_dict(self) -> dict: 

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

102 

103 Returns 

104 ------- 

105 result : 

106 The object encoded as a JSON compatible dict 

107 """ 

108 return { 

109 "origin": self.origin, 

110 "shape": self.model.shape, 

111 "peak": self.peak, 

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

113 "component_type": "component", 

114 "version": self.version, 

115 } 

116 

117 @classmethod 

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

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

120 

121 Parameters 

122 ---------- 

123 data : 

124 Dictionary representation of the object 

125 dtype : 

126 Datatype of the resulting model. 

127 

128 Returns 

129 ------- 

130 result : 

131 The reconstructed object 

132 """ 

133 data = MigrationRegistry.migrate(COMPONENT_TYPE, data) 

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

135 return cls( 

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

137 peak=data["peak"], 

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

139 ) 

140 

141 

142ScarletCubeComponentData.register() 

143 

144 

145@migration(COMPONENT_TYPE, PRE_SCHEMA) 

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

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

148 

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

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

151 

152 Parameters 

153 ---------- 

154 data : 

155 The data to migrate. 

156 

157 Returns 

158 ------- 

159 result : 

160 The migrated data. 

161 """ 

162 data["version"] = "1.0.0" 

163 return data