Coverage for tests/test_component.py: 18%

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 

25 

26import numpy as np 

27from lsst.scarlet.lite import Box, Image, Parameter 

28from lsst.scarlet.lite.component import ( 

29 Component, 

30 FactorizedComponent, 

31 default_adaprox_parameterization, 

32 default_fista_parameterization, 

33) 

34from lsst.scarlet.lite.operators import Monotonicity 

35from numpy.testing import assert_almost_equal, assert_array_equal 

36from utils import ScarletTestCase 

37 

38 

39class DummyComponent(Component): 

40 def resize(self) -> bool: 

41 pass 

42 

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

44 pass 

45 

46 def get_model(self) -> Image: 

47 pass 

48 

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

50 parameterization(self) 

51 

52 def to_data(self) -> DummyComponent: 

53 pass 

54 

55 

56class TestFactorizedComponent(ScarletTestCase): 

57 def setUp(self) -> None: 

58 spectrum = np.arange(3).astype(np.float32) 

59 morph = np.arange(20).reshape(4, 5).astype(np.float32) 

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

61 bbox = Box((4, 5), (22, 31)) 

62 self.model_box = Box((100, 100)) 

63 center = (24, 33) 

64 

65 self.component = FactorizedComponent( 

66 bands, 

67 spectrum, 

68 morph, 

69 bbox, 

70 center, 

71 ) 

72 

73 self.bands = bands 

74 self.spectrum = spectrum 

75 self.morph = morph 

76 self.full_shape = (3, 100, 100) 

77 

78 def test_constructor(self): 

79 # Test with only required parameters 

80 component = FactorizedComponent( 

81 self.bands, 

82 self.spectrum, 

83 self.morph, 

84 self.component.bbox, 

85 ) 

86 

87 self.assertIsInstance(component._spectrum, Parameter) 

88 assert_array_equal(component.spectrum, self.spectrum) 

89 self.assertIsInstance(component._morph, Parameter) 

90 assert_array_equal(component.morph, self.morph) 

91 self.assertBoxEqual(component.bbox, self.component.bbox) 

92 self.assertIsNone(component.peak) 

93 self.assertIsNone(component.bg_rms) 

94 self.assertEqual(component.bg_thresh, 0.25) 

95 self.assertEqual(component.floor, 1e-20) 

96 self.assertTupleEqual(component.shape, (3, 4, 5)) 

97 

98 # Test that parameters are passed through 

99 center = self.component.peak 

100 bg_rms = np.arange(5) / 10 

101 bg_thresh = 0.9 

102 floor = 1e-10 

103 

104 component = FactorizedComponent( 

105 self.bands, 

106 self.spectrum, 

107 self.morph, 

108 self.component.bbox, 

109 center, 

110 bg_rms, 

111 bg_thresh, 

112 floor, 

113 ) 

114 

115 self.assertTupleEqual(component.peak, center) 

116 assert_array_equal(component.bg_rms, bg_rms) # type: ignore 

117 self.assertEqual(component.bg_thresh, bg_thresh) 

118 self.assertEqual(component.floor, floor) 

119 self.assertEqual(component.get_model().dtype, np.float32) 

120 

121 def test_get_model(self): 

122 component = self.component 

123 assert_array_equal(component.get_model(), self.spectrum[:, None, None] * self.morph[None, :, :]) 

124 

125 # Insert component into a larger model 

126 full_model = np.zeros(self.full_shape) 

127 full_model[:, 22:26, 31:36] = self.spectrum[:, None, None] * self.morph[None, :, :] 

128 

129 test_model = Image(np.zeros(self.full_shape), bands=self.bands) 

130 test_model += component.get_model() 

131 

132 assert_array_equal(test_model.data, full_model) 

133 

134 def test_gradients(self): 

135 component = self.component 

136 morph = self.morph 

137 spectrum = self.spectrum 

138 

139 input_grad = np.array([morph, 2 * morph, 3 * morph]) 

140 true_spectrum_grad = np.array( 

141 [ 

142 np.sum(morph**2), 

143 np.sum(2 * morph**2), 

144 np.sum(3 * morph**2), 

145 ] 

146 ) 

147 assert_almost_equal(component.grad_spectrum(input_grad, spectrum, morph), true_spectrum_grad) 

148 

149 true_morph_grad = np.sum(input_grad * spectrum[:, None, None], axis=0) 

150 assert_almost_equal(component.grad_morph(input_grad, morph, spectrum), true_morph_grad) 

151 

152 def test_proximal_operators(self): 

153 # Test spectrum positivity, morph threshold, and monotonicity 

154 spectrum = np.array([-1, 2, 3], dtype=float) 

155 morph = np.array([[10, 2, 1], [1, 5, 3], [0.1, 4, -1]], dtype=float) 

156 bbox = Box((3, 3), (10, 10)) 

157 morph_bbox = Box((100, 100)) 

158 center = (11, 11) 

159 monotonicity = Monotonicity((101, 101), fit_radius=0) 

160 

161 component = FactorizedComponent( 

162 self.bands, 

163 spectrum.copy(), 

164 morph.copy(), 

165 bbox, 

166 center, 

167 bg_rms=np.array([1, 1, 1]), 

168 bg_thresh=0.5, 

169 monotonicity=monotonicity, 

170 ) 

171 

172 proxed_spectrum = np.array([1e-20, 2, 3]) 

173 proxed_morph = np.array([[2.6666666666666667, 2, 1], [1, 5, 3], [0, 4, 0]]) 

174 proxed_morph = proxed_morph / 5 

175 

176 component.prox_spectrum(component.spectrum) 

177 component.prox_morph(component.morph) 

178 

179 assert_array_equal(component.spectrum, proxed_spectrum) 

180 assert_array_equal(component.morph, proxed_morph) 

181 

182 component = FactorizedComponent( 

183 self.bands, 

184 spectrum.copy(), 

185 morph.copy(), 

186 bbox, 

187 None, 

188 ) 

189 

190 proxed_spectrum = np.array([1e-20, 2, 3]) 

191 proxed_morph = np.array([[10, 2, 1], [1, 5, 3], [0.1, 4, 0]]) 

192 proxed_morph = proxed_morph / 10 

193 

194 component.prox_spectrum(component.spectrum) 

195 component.prox_morph(component.morph) 

196 

197 assert_array_equal(component.spectrum, proxed_spectrum) 

198 assert_array_equal(component.morph, proxed_morph) 

199 

200 self.assertFalse(component.resize(morph_bbox)) 

201 

202 def test_resize(self): 

203 spectrum = np.array([1, 2, 3], dtype=float) 

204 morph = np.zeros((10, 10), dtype=float) 

205 morph[3:6, 5:8] = np.arange(9).reshape(3, 3) 

206 bbox = Box((10, 10), (3, 5)) 

207 

208 morph_bbox = Box((100, 100)) 

209 monotonicity = Monotonicity((101, 101), fit_radius=0) 

210 

211 component = FactorizedComponent( 

212 self.bands, 

213 spectrum.copy(), 

214 morph.copy(), 

215 bbox, 

216 None, 

217 bg_rms=np.array([1, 1, 1]), 

218 bg_thresh=0.5, 

219 monotonicity=monotonicity, 

220 padding=1, 

221 ) 

222 

223 self.assertTupleEqual(component.morph.shape, (10, 10)) 

224 self.assertIsNone(component.component_center) 

225 

226 component.resize(morph_bbox) 

227 self.assertTupleEqual(component.morph.shape, (5, 5)) 

228 self.assertTupleEqual(component.bbox.origin, (5, 9)) 

229 self.assertTupleEqual(component.bbox.shape, (5, 5)) 

230 self.assertIsNone(component.component_center) 

231 

232 def test_parameterization(self): 

233 component = self.component 

234 assert_array_equal(component.get_model(), self.spectrum[:, None, None] * self.morph[None, :, :]) 

235 

236 component.parameterize(default_fista_parameterization) 

237 helpers = set(component._morph.helpers.keys()) 

238 self.assertSetEqual(helpers, {"z"}) 

239 component.parameterize(default_adaprox_parameterization) 

240 helpers = set(component._morph.helpers.keys()) 

241 self.assertSetEqual(helpers, {"m", "v", "vhat"}) 

242 

243 params = (tuple("grizy"), Box((5, 5))) 

244 with self.assertRaises(NotImplementedError): 

245 default_fista_parameterization(DummyComponent(*params)) 

246 

247 with self.assertRaises(NotImplementedError): 

248 default_adaprox_parameterization(DummyComponent(*params))