Coverage for tests/test_io.py: 14%

1# This file is part of lsst.scarlet.lite. 

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 

22import json 

23import os 

24 

25import numpy as np 

26from lsst.scarlet.lite import Blend, Image, Observation, io 

27from lsst.scarlet.lite.initialization import FactorizedInitialization 

28from lsst.scarlet.lite.io import ComponentCube 

29from lsst.scarlet.lite.operators import Monotonicity 

30from lsst.scarlet.lite.utils import integrated_circular_gaussian 

31from numpy.testing import assert_almost_equal 

32from utils import ScarletTestCase 

33 

34 

35class TestIo(ScarletTestCase): 

36 def setUp(self) -> None: 

37 filename = os.path.join(__file__, "..", "..", "data", "hsc_cosmos_35.npz") 

38 filename = os.path.abspath(filename) 

39 data = np.load(filename) 

40 model_psf = integrated_circular_gaussian(sigma=0.8) 

41 self.detect = np.sum(data["images"], axis=0) 

42 self.centers = np.array([data["catalog"]["y"], data["catalog"]["x"]]).T 

43 bands = data["filters"] 

44 self.observation = Observation( 

45 Image(data["images"], bands=bands), 

46 Image(data["variance"], bands=bands), 

47 Image(1 / data["variance"], bands=bands), 

48 data["psfs"], 

49 model_psf[None], 

50 bands=bands, 

51 ) 

52 monotonicity = Monotonicity((101, 101)) 

53 init = FactorizedInitialization(self.observation, self.centers, monotonicity=monotonicity) 

54 self.blend = Blend(init.sources, self.observation) 

55 

56 def test_json(self): 

57 blend = self.blend 

58 blend.metadata = { 

59 "psf": self.observation.model_psf, 

60 "bands": tuple(str(band) for band in self.observation.bands), 

61 } 

62 blend_data = blend.to_data() 

63 metadata = { 

64 "model_psf": self.observation.model_psf, 

65 } 

66 model_data = io.ScarletModelData( 

67 blends={1: blend_data}, 

68 metadata=metadata, 

69 ) 

70 

71 # Get the json string for the model 

72 model_str = model_data.json() 

73 # Load the model string from the json 

74 model_dict = json.loads(model_str) 

75 # Load the full set of model data classes from the json string 

76 model_data = io.ScarletModelData.parse_obj(model_dict) 

77 metadata = model_data.metadata 

78 self.assertIsNotNone(metadata) 

79 # Convert the model data into scarlet models 

80 loaded_blend = model_data.blends[1].minimal_data_to_blend( 

81 model_psf=metadata["model_psf"], # type: ignore 

82 dtype=blend.observation.dtype, 

83 ) 

84 

85 self.assertEqual(len(blend.sources), len(loaded_blend.sources)) 

86 self.assertEqual(len(blend.components), len(loaded_blend.components)) 

87 self.assertImageAlmostEqual(blend.get_model(), loaded_blend.get_model()) 

88 

89 for sidx in range(len(blend.sources)): 

90 source1 = blend.sources[sidx] 

91 source2 = loaded_blend.sources[sidx] 

92 self.assertTupleEqual(source1.center, source2.center) 

93 self.assertEqual(len(source1.components), len(source2.components)) 

94 self.assertBoxEqual(source1.bbox, source2.bbox) 

95 for cidx in range(len(source1.components)): 

96 component1 = source1.components[cidx] 

97 component2 = source2.components[cidx] 

98 self.assertEqual(component1.peak, component2.peak) 

99 assert_almost_equal(component1.spectrum, component2.spectrum) 

100 assert_almost_equal(component1.morph, component2.morph) 

101 self.assertBoxEqual(component1.bbox, component2.bbox) 

102 

103 def test_cube_component(self): 

104 blend = self.blend 

105 for i in range(len(blend.sources)): 

106 blend.sources[i].metadata = {"id": f"peak-{i}"} 

107 component = blend.sources[-1].components[-1] 

108 # Replace one of the components with a Free-Form component. 

109 blend.sources[-1].components[-1] = ComponentCube( 

110 model=component.get_model(), 

111 peak=component.peak, 

112 ) 

113 

114 blend_data = blend.to_data() 

115 model_data = io.ScarletModelData( 

116 blends={1: blend_data}, 

117 metadata={ 

118 "model_psf": self.observation.model_psf, 

119 "psf": self.observation.psfs, 

120 "bands": tuple(str(band) for band in self.observation.bands), 

121 }, 

122 ) 

123 

124 # Get the json string for the model 

125 model_str = model_data.json() 

126 # Load the model string from the json 

127 model_dict = json.loads(model_str) 

128 # Load the full set of model data classes from the json string 

129 model_data = io.ScarletModelData.parse_obj(model_dict) 

130 # Convert the model data into scarlet models 

131 loaded_blend = model_data.blends[1].minimal_data_to_blend( 

132 model_psf=model_data.metadata["model_psf"], # type: ignore 

133 bands=model_data.metadata["bands"], # type: ignore 

134 psf=model_data.metadata["psf"], # type: ignore 

135 dtype=blend.observation.dtype, 

136 ) 

137 

138 self.assertEqual(len(blend.sources), len(loaded_blend.sources)) 

139 self.assertEqual(len(blend.components), len(loaded_blend.components)) 

140 self.assertImageAlmostEqual(blend.get_model(), loaded_blend.get_model()) 

141 

142 # Check that the metadata was stored correctly 

143 for i in range(len(blend.sources)): 

144 self.assertEqual(blend.sources[i].metadata, loaded_blend.sources[i].metadata) 

145 

146 def test_legacy_json(self): 

147 blend = self.blend 

148 

149 # Create legacy blend JSON data 

150 blend_data = blend.to_data().as_dict() 

151 encoded_psf = io.utils.numpy_to_json(self.observation.psfs) 

152 blend_data["psf"] = encoded_psf["data"] 

153 blend_data["psf_shape"] = encoded_psf["shape"] 

154 blend_data["bands"] = tuple(str(band) for band in self.observation.bands) 

155 blend_data["psf_center"] = (10, 10) 

156 

157 # Create legacy model data 

158 model_data = io.ScarletModelData(blends={}).as_dict() 

159 model_data["blends"][1] = blend_data 

160 encoded_psf = io.utils.numpy_to_json(self.observation.model_psf) 

161 model_data["psf"] = encoded_psf["data"] 

162 model_data["psfShape"] = encoded_psf["shape"] 

163 

164 # Legacy models were pre-versioning, so delete any version key 

165 model_data.pop("version", None) 

166 blend_data.pop("version", None) 

167 for source in blend_data["sources"].values(): 

168 source.pop("version", None) 

169 for component in source["components"]: 

170 component.pop("version", None) 

171 

172 self.assertIsNone(model_data["metadata"]) 

173 

174 # Get the json string for the model 

175 model_str = json.dumps(model_data) 

176 # Load the model string from the json 

177 model_dict = json.loads(model_str) 

178 # Load the full set of model data classes from the json string 

179 model_data = io.ScarletModelData.parse_obj(model_dict) 

180 metadata = model_data.metadata 

181 self.assertIsNotNone(metadata) 

182 

183 # Convert the model data into scarlet models 

184 loaded_blend = model_data.blends[1].minimal_data_to_blend( 

185 model_psf=metadata["model_psf"], # type: ignore 

186 dtype=blend.observation.dtype, 

187 ) 

188 

189 self.assertEqual(len(blend.sources), len(loaded_blend.sources)) 

190 self.assertEqual(len(blend.components), len(loaded_blend.components)) 

191 self.assertImageAlmostEqual(blend.get_model(), loaded_blend.get_model())