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

1from __future__ import annotations 

2 

3from dataclasses import dataclass 

4from typing import TYPE_CHECKING 

5 

6import numpy as np 

7from numpy.typing import DTypeLike 

8 

9from ..bbox import Box 

10from ..observation import Observation 

11from ..parameters import FixedParameter 

12from .component import ScarletComponentBaseData 

13from .migration import PRE_SCHEMA, MigrationRegistry, migration 

14 

15if TYPE_CHECKING: 

16 from ..component import FactorizedComponent 

17 

18__all__ = ["ScarletFactorizedComponentData"] 

19 

20CURRENT_SCHEMA = "1.0.0" 

21COMPONENT_TYPE = "factorized" 

22MigrationRegistry.set_current(COMPONENT_TYPE, CURRENT_SCHEMA) 

23 

24 

25@dataclass(kw_only=True) 

26class ScarletFactorizedComponentData(ScarletComponentBaseData): 

27 """Data for a factorized component 

28 

29 Attributes 

30 ---------- 

31 component_type : 

32 The type of component being stored. 

33 origin : 

34 The lower bound of the component's bounding box. 

35 peak : 

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

37 spectrum : 

38 The SED of the component. 

39 morph : 

40 The 2D morphology of the component. 

41 version : 

42 The schema version of the stored data. 

43 """ 

44 

45 component_type: str = COMPONENT_TYPE 

46 origin: tuple[int, int] 

47 peak: tuple[float, float] 

48 spectrum: np.ndarray 

49 morph: np.ndarray 

50 version: str = CURRENT_SCHEMA 

51 

52 @property 

53 def shape(self): 

54 """The shape of the component's morphology""" 

55 return self.morph.shape 

56 

57 def to_component(self, observation: Observation) -> FactorizedComponent: 

58 """Convert the storage data model into a scarlet FactorizedComponent 

59 

60 Parameters 

61 ---------- 

62 observation : 

63 The observation that the component is associated with 

64 

65 Returns 

66 ------- 

67 factorized_component : 

68 A scarlet factorized component extracted from persisted data. 

69 """ 

70 from ..component import FactorizedComponent 

71 

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

73 spectrum = self.spectrum 

74 morph = self.morph 

75 if self.peak is None: 

76 peak = None 

77 else: 

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

79 assert peak is not None 

80 # Note: since we aren't fitting a model, we don't need to 

81 # set the RMS of the background. 

82 # We set it to NaN just to be safe. 

83 component = FactorizedComponent( 

84 bands=observation.bands, 

85 spectrum=FixedParameter(spectrum), 

86 morph=FixedParameter(morph), 

87 peak=peak, 

88 bbox=bbox, 

89 bg_rms=np.full((len(observation.bands),), np.nan), 

90 ) 

91 return component 

92 

93 def as_dict(self) -> dict: 

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

95 

96 Returns 

97 ------- 

98 result : 

99 The object encoded as a JSON compatible dict 

100 """ 

101 return { 

102 "origin": tuple(int(o) for o in self.origin), 

103 "shape": tuple(int(s) for s in self.morph.shape), 

104 "peak": tuple(int(p) for p in self.peak), 

105 "spectrum": tuple(self.spectrum.astype(float)), 

106 "morph": tuple(self.morph.flatten().astype(float)), 

107 "component_type": self.component_type, 

108 "version": self.version, 

109 } 

110 

111 @classmethod 

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

113 """Reconstruct `ScarletFactorizedComponentData` from JSON compatible 

114 dict. 

115 

116 Parameters 

117 ---------- 

118 data : 

119 Dictionary representation of the object 

120 dtype : 

121 Datatype of the resulting model. 

122 

123 Returns 

124 ------- 

125 result : 

126 The reconstructed object 

127 """ 

128 data = MigrationRegistry.migrate(COMPONENT_TYPE, data) 

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

130 return cls( 

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

132 peak=data["peak"], 

133 spectrum=np.array(data["spectrum"]).astype(dtype), 

134 morph=np.array(data["morph"]).reshape(shape).astype(dtype), 

135 ) 

136 

137 

138ScarletFactorizedComponentData.register() 

139 

140 

141@migration(COMPONENT_TYPE, PRE_SCHEMA) 

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

143 """Migrate a pre-schema factorized component to schema version 1.0.0 

144 

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

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

147 

148 Parameters 

149 ---------- 

150 data : 

151 The data to migrate. 

152 

153 Returns 

154 ------- 

155 result : 

156 The migrated data. 

157 """ 

158 data["version"] = "1.0.0" 

159 return data