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 Attributes 

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 model_type : 

32 The type of model being stored. 

33 version : 

34 The schema version of the ScarletModelData. 

35 """ 

36 

37 model_type: str = MODEL_TYPE 

38 blends: dict[int, ScarletBlendBaseData] 

39 metadata: dict[str, Any] | None 

40 version: str = CURRENT_SCHEMA 

41 

42 def __init__( 

43 self, 

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

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

46 ): 

47 """Initialize an instance""" 

48 self.metadata = metadata 

49 if blends is None: 

50 blends = {} 

51 self.blends = blends 

52 

53 def as_dict(self) -> dict: 

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

55 

56 Returns 

57 ------- 

58 result : 

59 The object encoded as a JSON compatible dict 

60 """ 

61 result = { 

62 "model_type": self.model_type, 

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

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

65 "version": self.version, 

66 } 

67 return result 

68 

69 def json(self) -> str: 

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

71 

72 Returns 

73 ------- 

74 result : `str` 

75 The result of the object converted into a JSON format 

76 """ 

77 result = self.as_dict() 

78 return json.dumps(result) 

79 

80 @classmethod 

81 def from_dict( 

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

83 ) -> ScarletModelData: 

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

85 

86 Parameters 

87 ---------- 

88 data : 

89 Dictionary representation of the object 

90 dtype : 

91 Datatype of the resulting model. 

92 kwargs : 

93 Additional keyword arguments. 

94 

95 Returns 

96 ------- 

97 result : 

98 The reconstructed object 

99 """ 

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

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

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

103 if "blend_type" not in blend: 

104 # Assume that this is a legacy model 

105 blend["blend_type"] = "blend" 

106 try: 

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

108 except KeyError: 

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

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

111 

112 return cls( 

113 blends=blends, 

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

115 **kwargs, 

116 ) 

117 

118 @classmethod 

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

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

121 

122 Parameters 

123 ---------- 

124 data : 

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

126 

127 Returns 

128 ------- 

129 result : 

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

131 """ 

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

133 

134 

135@migration(MODEL_TYPE, PRE_SCHEMA) 

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

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

138 

139 Parameters 

140 ---------- 

141 data : 

142 The data to migrate. 

143 Returns 

144 ------- 

145 result : 

146 The migrated data. 

147 """ 

148 if "psfShape" in data: 

149 # Support legacy models before metadata was used 

150 data["metadata"] = { 

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

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

153 "array_keys": ["model_psf"], 

154 } 

155 data["version"] = "1.0.0" 

156 return data