Coverage for python/lsst/scarlet/lite/io/model_data.py: 45%

1from __future__ import annotations 

2 

3import json 

4from typing import Any 

5 

6import numpy as np 

7from numpy.typing import DTypeLike 

8 

9from .blend import ScarletBlendBaseData 

10from .migration import PRE_SCHEMA, MigrationRegistry, migration 

11from .utils import PersistenceError, decode_metadata, encode_metadata 

12 

13__all__ = ["ScarletModelData"] 

14 

15CURRENT_SCHEMA = "1.0.0" 

16MODEL_TYPE = "scarlet_model" 

17MigrationRegistry.set_current(MODEL_TYPE, CURRENT_SCHEMA) 

18 

19 

20class ScarletModelData: 

21 """A container that propagates scarlet models for an entire catalog. 

22 

23 Parameters 

24 ---------- 

25 blends : 

26 Map from parent IDs in the source catalog 

27 to scarlet model data for each parent ID (blend). 

28 metadata : 

29 Metadata associated with the model, 

30 for example the order of bands. 

31 

32 Attributes 

33 ---------- 

34 model_type : 

35 The type of model being stored. 

36 version : 

37 The schema version of the ScarletModelData. 

38 """ 

39 

40 model_type: str = MODEL_TYPE 

41 blends: dict[int, ScarletBlendBaseData] 

42 metadata: dict[str, Any] | None 

43 version: str = CURRENT_SCHEMA 

44 

45 def __init__( 

46 self, 

47 blends: dict[int, ScarletBlendBaseData] | None = None, 

48 metadata: dict[str, Any] | None = None, 

49 ): 

50 """Initialize an instance""" 

51 self.metadata = metadata 

52 if blends is None: 

53 blends = {} 

54 self.blends = blends 

55 

56 def as_dict(self) -> dict: 

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

58 

59 Returns 

60 ------- 

61 result : 

62 The object encoded as a JSON compatible dict 

63 """ 

64 result = { 

65 "model_type": self.model_type, 

66 "blends": {bid: blend_data.as_dict() for bid, blend_data in self.blends.items()}, 

67 "metadata": encode_metadata(self.metadata), 

68 "version": self.version, 

69 } 

70 return result 

71 

72 def json(self) -> str: 

73 """Serialize the data model to a JSON formatted string 

74 

75 Returns 

76 ------- 

77 result : `str` 

78 The result of the object converted into a JSON format 

79 """ 

80 result = self.as_dict() 

81 return json.dumps(result) 

82 

83 @classmethod 

84 def from_dict( 

85 cls, data: dict, dtype: DTypeLike = np.float32, **kwargs: dict[str, Any] 

86 ) -> ScarletModelData: 

87 """Reconstruct `ScarletModelData` from JSON compatible dict. 

88 

89 Parameters 

90 ---------- 

91 data : 

92 Dictionary representation of the object 

93 dtype : 

94 Datatype of the resulting model. 

95 kwargs : 

96 Additional keyword arguments. 

97 

98 Returns 

99 ------- 

100 result : 

101 The reconstructed object 

102 """ 

103 data = MigrationRegistry.migrate(cls.model_type, data) 

104 blends: dict[int, ScarletBlendBaseData] | None = {} 

105 for bid, blend in data.get("blends", {}).items(): 

106 if "blend_type" not in blend: 

107 # Assume that this is a legacy model 

108 blend["blend_type"] = "blend" 

109 try: 

110 blend_data = ScarletBlendBaseData.from_dict(blend, dtype=dtype) 

111 except KeyError: 

112 raise PersistenceError(f"Unknown blend type: {blend['blend_type']} for blend ID: {bid}") 

113 blends[int(bid)] = blend_data # type: ignore 

114 

115 return cls( 

116 blends=blends, 

117 metadata=decode_metadata(data["metadata"]), 

118 **kwargs, 

119 ) 

120 

121 @classmethod 

122 def parse_obj(cls, data: dict) -> ScarletModelData: 

123 """Construct a ScarletModelData from python decoded JSON object. 

124 

125 Parameters 

126 ---------- 

127 data : 

128 The result of json.load(s) on a JSON persisted ScarletModelData 

129 

130 Returns 

131 ------- 

132 result : 

133 The `ScarletModelData` that was loaded the from the input object 

134 """ 

135 return cls.from_dict(data, dtype=np.float32) 

136 

137 

138@migration(MODEL_TYPE, PRE_SCHEMA) 

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

140 """Migrate a pre-schema model to schema version 1.0.0 

141 

142 Parameters 

143 ---------- 

144 data : 

145 The data to migrate. 

146 Returns 

147 ------- 

148 result : 

149 The migrated data. 

150 """ 

151 if "psfShape" in data: 

152 # Support legacy models before metadata was used 

153 data["metadata"] = { 

154 "model_psf": data["psf"], 

155 "model_psf_shape": data["psfShape"], 

156 "array_keys": ["model_psf"], 

157 } 

158 data["version"] = "1.0.0" 

159 return data