Coverage for tests/test_blend.py: 21%

1# This file is part of 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 

22from __future__ import annotations 

23 

24from typing import Callable, cast 

25 

26import numpy as np 

27from lsst.scarlet.lite import Blend, Box, Image, Observation, Source 

28from lsst.scarlet.lite.component import Component, FactorizedComponent, default_adaprox_parameterization 

29from lsst.scarlet.lite.initialization import FactorizedInitialization 

30from lsst.scarlet.lite.operators import Monotonicity 

31from lsst.scarlet.lite.parameters import Parameter 

32from lsst.scarlet.lite.utils import integrated_circular_gaussian 

33from numpy.testing import assert_almost_equal, assert_raises 

34from scipy.signal import convolve as scipy_convolve 

35from utils import ObservationData, ScarletTestCase 

36 

37 

38class DummyCubeComponent(Component): 

39 def __init__(self, model: Image): 

40 super().__init__(model.bands, model.bbox) 

41 self._model = Parameter(model.data, {}, 0) 

42 

43 @property 

44 def data(self) -> np.ndarray: 

45 return self._model.x 

46 

47 def resize(self, model_box: Box) -> bool: 

48 pass 

49 

50 def update(self, it: int, input_grad: np.ndarray): 

51 pass 

52 

53 def get_model(self) -> Image: 

54 return Image(self.data, bands=self.bands, yx0=self.bbox.origin) 

55 

56 def parameterize(self, parameterization: Callable) -> None: 

57 pass 

58 

59 def to_data(self) -> DummyCubeComponent: 

60 pass 

61 

62 

63class TestBlend(ScarletTestCase): 

64 def setUp(self): 

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

66 yx0 = (1000, 2000) 

67 # The PSF in each band of the "observation" 

68 psfs = np.array([integrated_circular_gaussian(sigma=sigma) for sigma in [1.05, 0.9, 1.2]]) 

69 # The PSF of the model 

70 model_psf = integrated_circular_gaussian(sigma=0.8) 

71 

72 # The spectrum of each source 

73 spectra = np.array( 

74 [ 

75 [40, 10, 0], 

76 [0, 25, 40], 

77 [15, 8, 3], 

78 [20, 3, 4], 

79 [0, 30, 60], 

80 ], 

81 dtype=float, 

82 ) 

83 

84 # Use a point source for all of the sources 

85 morphs = [integrated_circular_gaussian(sigma=sigma) for sigma in [0.8, 2.5, 1.1, 2.1, 1.5]] 

86 # Make the second component a disk component 

87 morphs[1] = scipy_convolve(morphs[1], model_psf, mode="same") 

88 

89 # Give the first two components the same center, and unique centers 

90 # for the remaining sources 

91 centers = [ 

92 (1010, 2012), 

93 (1010, 2012), 

94 (1020, 2023), 

95 (1020, 2010), 

96 (1025, 2020), 

97 ] 

98 

99 # Create the simulated image and associated data products 

100 test_data = ObservationData(bands, psfs, spectra, morphs, centers, model_psf, yx0=yx0) 

101 

102 # Create the Observation 

103 variance = np.ones((3, 35, 35), dtype=float) * 1e-2 

104 weights = 1 / variance 

105 weights = weights / np.max(weights) 

106 self.observation = Observation( 

107 test_data.convolved, 

108 variance, 

109 weights, 

110 psfs, 

111 model_psf[None], 

112 bands=bands, 

113 bbox=Box(variance.shape[-2:], origin=yx0), 

114 ) 

115 self.data = test_data 

116 self.spectra = spectra 

117 self.centers = centers 

118 self.morphs = morphs 

119 

120 components = [] 

121 for spectrum, center, morph, data_morph in zip( 

122 self.spectra, self.centers, self.morphs, self.data.morphs 

123 ): 

124 components.append( 

125 FactorizedComponent( 

126 bands=bands, 

127 spectrum=spectrum, 

128 morph=morph, 

129 bbox=data_morph.bbox, 

130 peak=center, 

131 ) 

132 ) 

133 

134 sources = [Source(components[:2])] 

135 sources += [Source([component]) for component in components[2:]] 

136 

137 self.blend = Blend(sources, self.observation) 

138 

139 def test_exact(self): 

140 """Test that a blend model initialized with the exact solution 

141 builds the model correctly 

142 """ 

143 blend = self.blend 

144 self.assertEqual(len(blend.components), 5) 

145 self.assertEqual(len(blend.sources), 4) 

146 self.assertBoxEqual(blend.bbox, Box(self.data.images.shape[1:], self.observation.bbox.origin)) 

147 self.assertImageAlmostEqual(blend.get_model(), self.data.images) 

148 self.assertImageAlmostEqual(blend.get_model(convolve=True), self.observation.images) 

149 self.assertImageAlmostEqual( 

150 self.observation.convolve(blend.get_model(), mode="real"), 

151 self.observation.images, 

152 ) 

153 

154 # Test that the log likelihood is very small 

155 assert_almost_equal([blend.log_likelihood], [0]) 

156 

157 # Test that grad_log_likelihood updates the loss 

158 self.assertListEqual(blend.loss, []) 

159 blend._grad_log_likelihood() 

160 assert_almost_equal(blend.loss, [0]) 

161 

162 # Remove one of the sources and calculate the non-zero log_likelihood 

163 del blend.sources[-1] 

164 # Update the loss function and check that the loss changed 

165 blend._grad_log_likelihood() 

166 assert_almost_equal(blend.log_likelihood, -60.011720889007485) 

167 assert_almost_equal(blend.loss, [0, -60.011720889007485]) 

168 

169 def test_fit_spectra(self): 

170 """Test that fitting the spectra with exact morphologies is 

171 identical to the multiband image 

172 """ 

173 np.random.seed(0) 

174 blend = self.blend 

175 

176 # Change the initial spectra so that they can be fit later 

177 for component in blend.components: 

178 c = cast(FactorizedComponent, component) 

179 c.spectrum[:] = np.random.rand(3) * 10 

180 

181 with assert_raises(AssertionError): 

182 # Since the spectra have not yet been fit, 

183 # the model and images should not be equal 

184 self.assertImageEqual(blend.get_model(), self.data.images) 

185 

186 # We initialized all of the morphologies exactly, 

187 # so fitting the spectra should give a nearly exact solution 

188 blend.fit_spectra() 

189 

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

191 self.assertEqual(len(blend.sources), 4) 

192 self.assertBoxEqual(blend.bbox, self.observation.bbox) 

193 self.assertImageAlmostEqual(blend.get_model(), self.data.images) 

194 self.assertImageAlmostEqual(blend.get_model(convolve=True), self.observation.images) 

195 

196 def test_fit(self): 

197 observation = self.observation 

198 np.random.seed(0) 

199 images = observation.images.copy() 

200 noise = np.random.normal(size=observation.images.shape) * 1e-2 

201 observation.images._data += noise 

202 

203 monotonicity = Monotonicity((101, 101)) 

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

205 

206 blend = Blend(init.sources, self.observation).fit_spectra() 

207 blend.parameterize(default_adaprox_parameterization) 

208 blend.fit(100) 

209 

210 self.assertImageAlmostEqual(blend.get_model(convolve=True), images, decimal=1) 

211 

212 def test_non_factorized(self): 

213 np.random.seed(1) 

214 blend = self.blend 

215 # Remove the disk component from the first source 

216 model = self.spectra[1][:, None, None] * self.morphs[1][None, :, :] 

217 yx0 = blend.sources[0].components[1].bbox.origin 

218 blend.sources[0].components = blend.sources[0].components[:1] 

219 

220 # Change the initial spectra so that they can be fit later 

221 for component in blend.components: 

222 c = cast(FactorizedComponent, component) 

223 c.spectrum[:] = np.random.rand(3) * 10 

224 

225 with assert_raises(AssertionError): 

226 # Since the spectra have not yet been fit, 

227 # the model and images should not be equal 

228 self.assertImageEqual(blend.get_model(), self.data.images) 

229 

230 # Remove the disk component from the first source 

231 blend.sources[0].components = blend.sources[0].components[:1] 

232 # Create a new source for the disk with a non-factorized component 

233 component = DummyCubeComponent(Image(model, bands=self.blend.observation.bands, yx0=yx0)) 

234 blend.sources.append(Source([component])) 

235 

236 blend.fit_spectra() 

237 

238 self.assertEqual(len(blend.components), 5) 

239 self.assertEqual(len(blend.sources), 5) 

240 self.assertImageAlmostEqual(blend.get_model(), self.data.images) 

241 

242 def test_clipping(self): 

243 blend = self.blend 

244 

245 # Change the initial spectra so that they can be fit later 

246 for component in blend.components: 

247 c = cast(FactorizedComponent, component) 

248 c.spectrum[:] = np.random.rand(3) * 10 

249 

250 with assert_raises(AssertionError): 

251 # Since the spectra have not yet been fit, 

252 # the model and images should not be equal 

253 self.assertImageEqual(blend.get_model(), self.data.images) 

254 

255 # Add an empty source 

256 zero_model = Image.from_box(Box((5, 5), (30, 0)), bands=blend.observation.bands) 

257 component = DummyCubeComponent(zero_model) 

258 blend.sources.append(Source([component])) 

259 

260 blend.fit_spectra(clip=True) 

261 

262 self.assertEqual(len(blend.components), 5) 

263 self.assertEqual(len(blend.sources), 5) 

264 self.assertImageAlmostEqual(blend.get_model(), self.data.images)