Coverage for python/lsst/meas/extensions/scarlet/io/model_data.py: 95%

1# This file is part of meas_extensions_scarlet. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24from typing import Any 

25 

26import numpy as np 

27from numpy.typing import DTypeLike 

28 

29import lsst.scarlet.lite as scl 

30 

31from .source_data import IsolatedSourceData 

32 

33__all__ = ["LsstScarletModelData"] 

34 

35CURRENT_SCHEMA = "1.0.1" 

36SCARLET_LITE_SCHEMA = "1.0.0" 

37MODEL_TYPE = "lsst" 

38scl.io.migration.MigrationRegistry.set_current(MODEL_TYPE, CURRENT_SCHEMA) 

39 

40 

41def _checkScarletLiteSchema(scarletSchema: str, pinnedSchema: str) -> None: 

42 """Raise if the installed scarlet_lite schema is not the schema 

43 this package was last verified against. 

44 

45 Bidirectional drift guard. Any mismatch means the IO layer 

46 cannot be trusted to round-trip data: an older installed 

47 scarlet may not emit the keys this package expects, and a 

48 newer one may have changed them. Either way the user gets an 

49 actionable error at import time instead of a confusing failure 

50 deep inside a ``from_dict`` call. 

51 

52 Parameters 

53 ---------- 

54 scarletSchema : str 

55 Schema string from the installed 

56 ``scl.io.model_data.CURRENT_SCHEMA``. 

57 pinnedSchema : str 

58 Schema string this package was last written against 

59 (``SCARLET_LITE_SCHEMA``). 

60 

61 Raises 

62 ------ 

63 RuntimeError 

64 If ``scarletSchema`` differs from ``pinnedSchema``. 

65 """ 

66 if scarletSchema != pinnedSchema: 

67 raise RuntimeError( 

68 "Version mismatch between meas_extensions_scarlet and scarlet lite. " 

69 "This requires updating SCARLET_LITE_SCHEMA, CURRENT_SCHEMA, and a migration step " 

70 f"to match the ScarletModelData schema version {scarletSchema}." 

71 ) 

72 

73 

74# Ensure that the ScarletModelData from scarlet lite hasn't changed. 

75_checkScarletLiteSchema(scl.io.model_data.CURRENT_SCHEMA, SCARLET_LITE_SCHEMA) 

76 

77 

78class LsstScarletModelData(scl.io.ScarletModelData): 

79 """A ScarletModelData that includes isolated sources. 

80 

81 Attributes 

82 ---------- 

83 isolated : dict[int, IsolatedSourceData] 

84 A mapping of isolated source IDs to their data. 

85 version : dict[int, scl.io.ScarletBlendBaseData] 

86 The schema version of the serialized data. 

87 """ 

88 model_type: str = MODEL_TYPE 

89 isolated: dict[int, IsolatedSourceData] 

90 version: str = CURRENT_SCHEMA 

91 

92 def __init__( 

93 self, 

94 isolated: dict[int, IsolatedSourceData] | None = None, 

95 blends: dict[int, scl.io.ScarletBlendBaseData] | None = None, 

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

97 ): 

98 super().__init__(blends=blends, metadata=metadata) 

99 self.isolated = isolated if isolated is not None else {} 

100 

101 def as_dict(self) -> dict[str, Any]: 

102 """Convert to a dictionary for serialization 

103 

104 Returns 

105 ------- 

106 result : dict[str, Any] 

107 The object encoded as a JSON-compatible dictionary. 

108 """ 

109 data = super().as_dict() 

110 data.update( 

111 { 

112 "model_type": MODEL_TYPE, 

113 "isolated": {k: v.as_dict() for k, v in self.isolated.items()}, 

114 "version": self.version, 

115 } 

116 ) 

117 return data 

118 

119 @classmethod 

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

121 """Reconstruct `LsstScarletModelData` from JSON compatible dict. 

122 

123 Parameters 

124 ---------- 

125 data : dict 

126 Dictionary representation of the object 

127 dtype : DTypeLike 

128 Datatype of the resulting model. 

129 

130 Returns 

131 ------- 

132 result : LsstScarletModelData 

133 The reconstructed object 

134 """ 

135 data = scl.io.migration.MigrationRegistry.migrate(MODEL_TYPE, data) 

136 isolated: dict[int, IsolatedSourceData] = {} 

137 for sid, source_data in data.get("isolated", {}).items(): 

138 isolated[int(sid)] = IsolatedSourceData.from_dict(source_data, dtype=dtype) 

139 if "metadata" not in data: 139 ↛ 140line 139 didn't jump to line 140 because the condition on line 139 was never true

140 data["metadata"] = None 

141 return super().from_dict(data, dtype=dtype, isolated=isolated) 

142 

143 

144@scl.io.migration.migration(MODEL_TYPE, scl.io.migration.PRE_SCHEMA) 

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

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

147 

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

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

150 

151 Parameters 

152 ---------- 

153 data : dict 

154 The data to migrate. 

155 Returns 

156 ------- 

157 result : dict 

158 The migrated data. 

159 """ 

160 # Ensure that the model type and version are set and add an 

161 # empty isolated sources dictionary. 

162 if "model_type" not in data: 162 ↛ 164line 162 didn't jump to line 164 because the condition on line 162 was always true

163 data["model_type"] = MODEL_TYPE 

164 data["isolated"] = {} 

165 # Pre-``metadata`` archives stored the model PSF as top-level 

166 # ``psf`` / ``psfShape`` entries. Promote them into the modern 

167 # ``metadata`` shape so ``decode_metadata`` can reconstruct the 

168 # array via ``array_keys``. Mirrors scarlet_lite's pre-schema 

169 # ``scarlet_model`` migration. 

170 if "metadata" not in data and "psfShape" in data: 

171 data["metadata"] = { 

172 "model_psf": data.pop("psf"), 

173 "model_psf_shape": data.pop("psfShape"), 

174 "array_keys": ["model_psf"], 

175 } 

176 data["version"] = "1.0.0" 

177 return data 

178 

179 

180@scl.io.migration.migration(MODEL_TYPE, "1.0.0") 

181def _to_1_0_1(data: dict) -> dict: 

182 """Migrate a schema version 1.0.0 model to schema version 1.0.1 

183 

184 There were no changes to this data model in v1.0.1 but we need 

185 to provide a way to migrate 1.0.0 data. 

186 

187 Parameters 

188 ---------- 

189 data : dict 

190 The data to migrate. 

191 Returns 

192 ------- 

193 result : dict 

194 The migrated data. 

195 """ 

196 data["version"] = "1.0.1" 

197 if data.get("metadata") is None: 

198 data["metadata"] = {} 

199 data["metadata"].setdefault("footprint", None) 

200 return data