Coverage for tests/test_io.py: 100%

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, Box, Image, Observation, io 

27from lsst.scarlet.lite.component import CubeComponent 

28from lsst.scarlet.lite.initialization import FactorizedInitialization 

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 self.assertBoxEqual(blend.bbox, blend_data.bbox) 

89 

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

91 source1 = blend.sources[sidx] 

92 source2 = loaded_blend.sources[sidx] 

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

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

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

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

97 component1 = source1.components[cidx] 

98 component2 = source2.components[cidx] 

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

100 assert_almost_equal(component1.spectrum, component2.spectrum) 

101 assert_almost_equal(component1.morph, component2.morph) 

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

103 

104 def test_cube_component(self): 

105 blend = self.blend 

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

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

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

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

110 blend.sources[-1].components[-1] = CubeComponent( 

111 model=component.get_model(), 

112 peak=component.peak, 

113 ) 

114 

115 blend_data = blend.to_data() 

116 model_data = io.ScarletModelData( 

117 blends={1: blend_data}, 

118 metadata={ 

119 "model_psf": self.observation.model_psf, 

120 "psf": self.observation.psfs, 

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

122 }, 

123 ) 

124 

125 # Get the json string for the model 

126 model_str = model_data.json() 

127 # Load the model string from the json 

128 model_dict = json.loads(model_str) 

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

130 model_data = io.ScarletModelData.parse_obj(model_dict) 

131 # Convert the model data into scarlet models 

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

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

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

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

136 dtype=blend.observation.dtype, 

137 ) 

138 

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

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

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

142 

143 # Check that the metadata was stored correctly 

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

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

146 

147 def test_cube_component_to_component_preserves_peak(self): 

148 """``ScarletCubeComponentData.to_component`` must preserve the 

149 full ``(y, x)`` peak, not collapse both axes onto ``peak[0]``. 

150 

151 Regression test: the implementation previously read ``peak[0]`` 

152 twice when constructing the returned ``CubeComponent``, so any 

153 non-symmetric peak silently round-tripped as ``(y, y)``. 

154 """ 

155 peak = (54, 105) 

156 n_bands, h, w = 3, 8, 10 

157 cube_data = io.ScarletCubeComponentData( 

158 origin=(50, 100), 

159 peak=peak, 

160 model=np.zeros((n_bands, h, w), dtype=np.float32), 

161 ) 

162 observation = Observation.empty( 

163 bands=("g", "r", "i"), 

164 psfs=np.ones((n_bands, 5, 5), dtype=np.float32), 

165 model_psf=np.ones((1, 5, 5), dtype=np.float32), 

166 bbox=Box((h, w), origin=(50, 100)), 

167 dtype=np.float32, 

168 ) 

169 

170 component = cube_data.to_component(observation) 

171 

172 self.assertEqual(component.peak, peak) 

173 

174 def test_legacy_json(self): 

175 blend = self.blend 

176 

177 # Create legacy blend JSON data 

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

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

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

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

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

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

184 

185 # Create legacy model data 

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

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

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

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

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

191 

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

193 model_data.pop("version", None) 

194 blend_data.pop("version", None) 

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

196 source.pop("version", None) 

197 for component in source["components"]: 

198 component.pop("version", None) 

199 

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

201 

202 # Get the json string for the model 

203 model_str = json.dumps(model_data) 

204 # Load the model string from the json 

205 model_dict = json.loads(model_str) 

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

207 model_data = io.ScarletModelData.parse_obj(model_dict) 

208 metadata = model_data.metadata 

209 self.assertIsNotNone(metadata) 

210 

211 # Convert the model data into scarlet models 

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

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

214 dtype=blend.observation.dtype, 

215 ) 

216 

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

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

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