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

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 

33CURRENT_SCHEMA = "1.0.0" 

34SCARLET_LITE_SCHEMA = "1.0.0" 

35MODEL_TYPE = "lsst" 

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

37 

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

39if scl.io.model_data.CURRENT_SCHEMA != SCARLET_LITE_SCHEMA: 39 ↛ 40line 39 didn't jump to line 40 because the condition on line 39 was never true

40 raise RuntimeError( 

41 "Version mismatch between meas_extensions_scarlet and scarlet lite. " 

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

43 f"to match the ScarletModelData schema version {scl.io.model_data.CURRENT_SCHEMA}." 

44 ) 

45 

46 

47class LsstScarletModelData(scl.io.ScarletModelData): 

48 """A ScarletModelData that includes isolated sources. 

49 

50 Attributes 

51 ---------- 

52 isolated : dict[int, IsolatedSourceData] 

53 A mapping of isolated source IDs to their data. 

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

55 The schema version of the serialized data. 

56 """ 

57 isolated: dict[int, IsolatedSourceData] 

58 version: str = CURRENT_SCHEMA 

59 

60 def __init__( 

61 self, 

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

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

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

65 ): 

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

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

68 

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

70 """Convert to a dictionary for serialization 

71 

72 Returns 

73 ------- 

74 result : dict[str, Any] 

75 The object encoded as a JSON-compatible dictionary. 

76 """ 

77 data = super().as_dict() 

78 # Ensure that the we are not trying to serialize with an 

79 # incompatible schema. 

80 if data["version"] != SCARLET_LITE_SCHEMA: 

81 raise RuntimeError( 

82 f"Cannot serialize scarlet model with schema {data['version']}, " 

83 f"expected {SCARLET_LITE_SCHEMA}. " 

84 "This requires updating SCARLET_LITE_SCHEMA, CURRENT_SCHEMA, and a migration step." 

85 ) 

86 data.update( 

87 { 

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

89 "version": self.version, 

90 } 

91 ) 

92 return data 

93 

94 @classmethod 

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

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

97 

98 Parameters 

99 ---------- 

100 data : dict 

101 Dictionary representation of the object 

102 dtype : DTypeLike 

103 Datatype of the resulting model. 

104 

105 Returns 

106 ------- 

107 result : LsstScarletModelData 

108 The reconstructed object 

109 """ 

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

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

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

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

114 if "metadata" not in data: 

115 data["metadata"] = None 

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

117 

118 

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

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

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

122 

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

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

125 

126 Parameters 

127 ---------- 

128 data : dict 

129 The data to migrate. 

130 Returns 

131 ------- 

132 result : dict 

133 The migrated data. 

134 """ 

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

136 # empty isolated sources dictionary. 

137 if "model_type" not in data: 

138 data["model_type"] = MODEL_TYPE 

139 data["isolated"] = {} 

140 data["version"] = "1.0.0" 

141 return data