Coverage for tests/test_blend.py: 19%

125 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-19 10:38 +0000

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 typing import Callable, cast 

23 

24import numpy as np 

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

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

27from lsst.scarlet.lite.initialization import FactorizedChi2Initialization 

28from lsst.scarlet.lite.operators import Monotonicity 

29from lsst.scarlet.lite.parameters import Parameter 

30from lsst.scarlet.lite.utils import integrated_circular_gaussian 

31from numpy.testing import assert_almost_equal, assert_raises 

32from scipy.signal import convolve as scipy_convolve 

33from utils import ObservationData, ScarletTestCase 

34 

35 

36class DummyCubeComponent(Component): 

37 def __init__(self, model: Image): 

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

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

40 

41 @property 

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

43 return self._model.x 

44 

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

46 pass 

47 

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

49 pass 

50 

51 def get_model(self) -> Image: 

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

53 

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

55 pass 

56 

57 

58class TestBlend(ScarletTestCase): 

59 def setUp(self): 

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

61 yx0 = (1000, 2000) 

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

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

64 # The PSF of the model 

65 model_psf = integrated_circular_gaussian(sigma=0.8) 

66 

67 # The spectrum of each source 

68 spectra = np.array( 

69 [ 

70 [40, 10, 0], 

71 [0, 25, 40], 

72 [15, 8, 3], 

73 [20, 3, 4], 

74 [0, 30, 60], 

75 ], 

76 dtype=float, 

77 ) 

78 

79 # Use a point source for all of the sources 

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

81 # Make the second component a disk component 

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

83 

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

85 # for the remaining sources 

86 centers = [ 

87 (1010, 2012), 

88 (1010, 2012), 

89 (1020, 2023), 

90 (1020, 2010), 

91 (1025, 2020), 

92 ] 

93 

94 # Create the simulated image and associated data products 

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

96 

97 # Create the Observation 

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

99 weights = 1 / variance 

100 weights = weights / np.max(weights) 

101 self.observation = Observation( 

102 test_data.convolved, 

103 variance, 

104 weights, 

105 psfs, 

106 model_psf[None], 

107 bands=bands, 

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

109 ) 

110 self.data = test_data 

111 self.spectra = spectra 

112 self.centers = centers 

113 self.morphs = morphs 

114 

115 components = [] 

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

117 self.spectra, self.centers, self.morphs, self.data.morphs 

118 ): 

119 components.append( 

120 FactorizedComponent( 

121 bands=bands, 

122 spectrum=spectrum, 

123 morph=morph, 

124 bbox=data_morph.bbox, 

125 peak=center, 

126 ) 

127 ) 

128 

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

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

131 

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

133 

134 def test_exact(self): 

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

136 builds the model correctly 

137 """ 

138 blend = self.blend 

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

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

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

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

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

144 self.assertImageAlmostEqual( 

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

146 self.observation.images, 

147 ) 

148 

149 # Test that the log likelihood is very small 

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

151 

152 # Test that grad_log_likelihood updates the loss 

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

154 blend._grad_log_likelihood() 

155 assert_almost_equal(blend.loss, [0]) 

156 

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

158 del blend.sources[-1] 

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

160 blend._grad_log_likelihood() 

161 assert_almost_equal(blend.log_likelihood, -60.011720889007485) 

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

163 

164 def test_fit_spectra(self): 

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

166 identical to the multiband image 

167 """ 

168 np.random.seed(0) 

169 blend = self.blend 

170 

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

172 for component in blend.components: 

173 c = cast(FactorizedComponent, component) 

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

175 

176 with assert_raises(AssertionError): 

177 # Since the spectra have not yet been fit, 

178 # the model and images should not be equal 

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

180 

181 # We initialized all of the morphologies exactly, 

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

183 blend.fit_spectra() 

184 

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

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

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

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

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

190 

191 def test_fit(self): 

192 observation = self.observation 

193 np.random.seed(0) 

194 images = observation.images.copy() 

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

196 observation.images._data += noise 

197 

198 monotonicity = Monotonicity((101, 101)) 

199 init = FactorizedChi2Initialization(observation, self.centers, monotonicity=monotonicity) 

200 

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

202 blend.parameterize(default_adaprox_parameterization) 

203 blend.fit(100) 

204 

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

206 

207 def test_non_factorized(self): 

208 np.random.seed(1) 

209 blend = self.blend 

210 # Remove the disk component from the first source 

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

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

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

214 

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

216 for component in blend.components: 

217 c = cast(FactorizedComponent, component) 

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

219 

220 with assert_raises(AssertionError): 

221 # Since the spectra have not yet been fit, 

222 # the model and images should not be equal 

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

224 

225 # Remove the disk component from the first source 

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

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

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

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

230 

231 blend.fit_spectra() 

232 

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

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

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

236 

237 def test_clipping(self): 

238 blend = self.blend 

239 

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

241 for component in blend.components: 

242 c = cast(FactorizedComponent, component) 

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

244 

245 with assert_raises(AssertionError): 

246 # Since the spectra have not yet been fit, 

247 # the model and images should not be equal 

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

249 

250 # Add an empty source 

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

252 component = DummyCubeComponent(zero_model) 

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

254 

255 blend.fit_spectra(clip=True) 

256 

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

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

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