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 def __init__(self, model: Image, peak: tuple[int, int]): 

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

37 

38 

39@dataclass(kw_only=True) 

40class ScarletCubeComponentData(ScarletComponentBaseData): 

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

42 

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

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

45 

46 Attributes 

47 ---------- 

48 origin : 

49 The lower bound of the components bounding box. 

50 peak : 

51 The peak of the component. 

52 model : 

53 The model for the component. 

54 """ 

55 

56 origin: tuple[int, int] 

57 peak: tuple[float, float] 

58 model: np.ndarray 

59 component_type: str = COMPONENT_TYPE 

60 version: str = CURRENT_SCHEMA 

61 

62 @property 

63 def shape(self): 

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

65 

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

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

68 

69 Parameters 

70 ---------- 

71 observation : 

72 The observation that the component is associated with 

73 

74 Returns 

75 ------- 

76 component : 

77 A scarlet component extracted from persisted data. 

78 """ 

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

80 model = self.model 

81 if self.peak is None: 

82 peak = None 

83 else: 

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

85 assert peak is not None 

86 component = CubeComponent( 

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

88 peak=peak, 

89 ) 

90 return component 

91 

92 def as_dict(self) -> dict: 

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

94 

95 Returns 

96 ------- 

97 result : 

98 The object encoded as a JSON compatible dict 

99 """ 

100 return { 

101 "origin": self.origin, 

102 "shape": self.model.shape, 

103 "peak": self.peak, 

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

105 "component_type": "component", 

106 "version": self.version, 

107 } 

108 

109 @classmethod 

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

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

112 

113 Parameters 

114 ---------- 

115 data : 

116 Dictionary representation of the object 

117 dtype : 

118 Datatype of the resulting model. 

119 

120 Returns 

121 ------- 

122 result : 

123 The reconstructed object 

124 """ 

125 data = MigrationRegistry.migrate(COMPONENT_TYPE, data) 

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

127 return cls( 

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

129 peak=data["peak"], 

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

131 ) 

132 

133 

134ScarletCubeComponentData.register() 

135 

136 

137@migration(COMPONENT_TYPE, PRE_SCHEMA) 

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

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

140 

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

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

143 

144 Parameters 

145 ---------- 

146 data : 

147 The data to migrate. 

148 

149 Returns 

150 ------- 

151 result : 

152 The migrated data. 

153 """ 

154 data["version"] = "1.0.0" 

155 return data