Coverage for tests/test_observation.py: 13%

103 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-20 03:40 -0700

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 numpy as np 

23from lsst.scarlet.lite import Box, Image, Observation 

24from lsst.scarlet.lite.observation import convolve as scarlet_convolve 

25from lsst.scarlet.lite.observation import get_filter_bounds, get_filter_coords 

26from lsst.scarlet.lite.utils import integrated_circular_gaussian 

27from numpy.testing import assert_almost_equal, assert_array_equal 

28from scipy.signal import convolve as scipy_convolve 

29from utils import ObservationData, ScarletTestCase 

30 

31 

32class TestObservation(ScarletTestCase): 

33 def setUp(self): 

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

35 variance = np.ones((3, 35, 35), dtype=float) 

36 weights = 1 / variance 

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

38 model_psf = integrated_circular_gaussian(sigma=0.8) 

39 

40 # The spectrum of each source 

41 spectra = np.array( 

42 [ 

43 [31, 10, 0], 

44 [0, 5, 20], 

45 [15, 8, 3], 

46 [20, 3, 4], 

47 [0, 30, 60], 

48 ] 

49 ) 

50 

51 # Use a point source for all of the sources 

52 morphs = [integrated_circular_gaussian(sigma=sigma) for sigma in [0.8, 3.1, 1.1, 2.1, 1.5]] 

53 # Make the second component a disk component 

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

55 

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

57 # for the remaining sources 

58 centers = [ 

59 (10, 12), 

60 (10, 12), 

61 (20, 23), 

62 (20, 10), 

63 (25, 20), 

64 ] 

65 

66 # Create the Observation 

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

68 self.observation = Observation( 

69 test_data.convolved, 

70 variance, 

71 weights, 

72 psfs, 

73 model_psf[None], 

74 bands=bands, 

75 ) 

76 self.data = test_data 

77 self.spectra = spectra 

78 self.centers = centers 

79 self.morphs = morphs 

80 self.psfs = psfs 

81 self.bands = bands 

82 

83 def tearDown(self): 

84 del self.data 

85 

86 def test_real_convolution_function(self): 

87 """Test that real space convolution works""" 

88 images = self.data.images 

89 true_convolved = np.array( 

90 [ 

91 scipy_convolve(images.data[b], self.observation.psfs[b], mode="same") 

92 for b in range(len(images.data)) 

93 ] 

94 ) 

95 coords = get_filter_coords(self.observation.psfs[0]) 

96 bounds = get_filter_bounds(coords.reshape(-1, 2)) 

97 convolved = scarlet_convolve(images.data, self.observation.psfs, bounds) 

98 assert_almost_equal(convolved, true_convolved) 

99 

100 with self.assertRaises(ValueError): 

101 get_filter_coords(np.arange(10)) 

102 

103 with self.assertRaises(ValueError): 

104 get_filter_coords(np.arange(16).reshape(4, 4)) 

105 

106 def test_constructors(self): 

107 np.random.seed(1) 

108 variance = np.random.normal(size=self.data.convolved.shape) ** 2 

109 observation = Observation( 

110 self.data.convolved, 

111 variance, 

112 1 / variance, 

113 self.psfs, 

114 bands=self.bands, 

115 ) 

116 self.assertImageEqual(observation.images, self.data.convolved) 

117 self.assertImageEqual(observation.variance, Image(variance, bands=self.bands)) 

118 self.assertImageEqual(observation.weights, Image(1 / variance, bands=self.bands)) 

119 assert_array_equal(observation.psfs, self.psfs) 

120 self.assertIsNone(observation.model_psf) 

121 self.assertIsNone(observation.diff_kernel) 

122 self.assertIsNone(observation.grad_kernel) 

123 assert_array_equal(observation.noise_rms, np.mean(np.sqrt(variance), axis=(1, 2))) 

124 self.assertBoxEqual(observation.bbox, Box(variance.shape[-2:])) 

125 self.assertIn(observation.mode, ["fft", "real"]) 

126 

127 # Set all of the pixels in the model to 1 more than the images, 

128 # so that images - model = np.ones, meaning the log_likelihood 

129 # is just half the sum of the weights. 

130 model = self.observation.images - 1 

131 self.assertEqual(observation.log_likelihood(model), -0.5 * np.sum(observation.weights.data)) 

132 self.assertTupleEqual(observation.shape, (3, 35, 35)) 

133 self.assertEqual(observation.n_bands, 3) 

134 self.assertEqual(observation.dtype, float) 

135 

136 def test_convolve(self): 

137 # Test the default initialization with no model psf, 

138 # menaing observation.convolve is a pass-through operation. 

139 np.random.seed(1) 

140 variance = np.random.normal(size=self.data.convolved.shape) ** 2 

141 observation = Observation( 

142 self.data.convolved, 

143 variance, 

144 1 / variance, 

145 self.psfs, 

146 bands=self.bands, 

147 ) 

148 assert_array_equal(observation.convolve(observation.images), observation.images) 

149 

150 # Use an observation with a model_psf and difference kernel and check 

151 # convolution. 

152 observation = self.observation 

153 assert_array_equal(observation.diff_kernel.image, self.data.diff_kernel.image) 

154 assert_almost_equal(observation.convolve(self.data.images).data, observation.images.data) 

155 

156 # Test real conversions 

157 deconvolved = self.data.images 

158 observation.mode = "real" 

159 self.assertImageAlmostEqual(observation.convolve(deconvolved), observation.images) 

160 

161 # Test convolution with the gradient 

162 grad_convolved = np.array( 

163 [ 

164 scipy_convolve( 

165 deconvolved.data[band], 

166 observation.grad_kernel.image[band], 

167 mode="same", 

168 ) 

169 for band in range(len(deconvolved.data)) 

170 ] 

171 ) 

172 assert_almost_equal(observation.convolve(deconvolved, grad=True).data, grad_convolved) 

173 

174 # Test that overriding the mode works 

175 real = observation.convolve(deconvolved, mode="real") 

176 self.assertImageAlmostEqual(real, observation.images) 

177 

178 with self.assertRaises(ValueError): 

179 observation.convolve(deconvolved, mode="fake") 

180 

181 # Test convolving a small image 

182 x = np.linspace(-3, 3, 7) 

183 small_array = integrated_circular_gaussian(x=x, y=x, sigma=0.8) 

184 small_psf = Image(np.array([small_array, small_array, small_array]), bands=observation.bands) 

185 truth = Image( 

186 np.array( 

187 [ 

188 scipy_convolve( 

189 small_psf[band].data, 

190 observation.diff_kernel.image[observation.bands.index(band)], 

191 method="direct", 

192 mode="same", 

193 ) 

194 for band in observation.bands 

195 ] 

196 ), 

197 bands=observation.bands, 

198 ) 

199 convolved = observation.convolve(small_psf) 

200 self.assertImageAlmostEqual(convolved, truth) 

201 

202 def test_index_extraction(self): 

203 alpha_bands = ("g", "i", "r", "y", "z") 

204 images = np.arange(60).reshape(5, 3, 4) 

205 image_g = images[0] 

206 image_i = images[1] 

207 image_r = images[2] 

208 variance = np.arange(5)[:, None, None] * np.ones((5, 3, 4)) + 1 

209 weights = 1 / variance 

210 psfs = np.arange(5)[:, None, None] * np.ones((5, 2, 2)) 

211 model_psf = np.zeros(9).reshape(3, 3) 

212 model_psf[1] = 1 

213 model_psf[:, 1] = 1 

214 model_psf[1, 1] = 2 

215 observation = Observation( 

216 images, 

217 variance, 

218 weights, 

219 psfs, 

220 model_psf[None], 

221 bands=alpha_bands, 

222 ) 

223 

224 bands = "i" 

225 truth = Image( 

226 np.arange(12, 24).reshape(3, 4), 

227 ) 

228 self.assertImageEqual(observation.images[bands], truth) 

229 

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

231 indices = observation.images.spectral_indices(bands) 

232 assert_array_equal(indices, (0, 2, 1)) 

233 self.assertImageEqual( 

234 observation.images[bands], 

235 Image(np.array([image_g, image_r, image_i]), bands=("g", "r", "i")), 

236 )