Coverage for tests/test_models.py: 13%

191 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-07 11:26 +0000

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 os 

23from functools import partial 

24from typing import cast 

25 

26import lsst.scarlet.lite.models as models 

27import numpy as np 

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

29from lsst.scarlet.lite.component import ( 

30 Component, 

31 default_adaprox_parameterization, 

32 default_fista_parameterization, 

33) 

34from lsst.scarlet.lite.initialization import FactorizedChi2Initialization 

35from lsst.scarlet.lite.models import ( 

36 CartesianFrame, 

37 EllipseFrame, 

38 EllipticalParametricComponent, 

39 FittedPsfBlend, 

40 FittedPsfObservation, 

41 FreeFormComponent, 

42 ParametricComponent, 

43) 

44from lsst.scarlet.lite.operators import Monotonicity 

45from lsst.scarlet.lite.parameters import AdaproxParameter, parameter, relative_step 

46from lsst.scarlet.lite.utils import integrated_circular_gaussian 

47from numpy.testing import assert_array_equal 

48from utils import ScarletTestCase 

49 

50 

51def parameterize(component: Component): 

52 assert isinstance(component, ParametricComponent) 

53 component._spectrum = AdaproxParameter( 

54 component.spectrum, 

55 step=partial(relative_step, factor=1e-2, minimum=1e-16), 

56 ) 

57 component._params = AdaproxParameter( 

58 component._params.x, 

59 step=1e-2, 

60 ) 

61 

62 

63class TestFreeForm(ScarletTestCase): 

64 def setUp(self) -> None: 

65 yx0 = (1000, 2000) 

66 filename = os.path.join(__file__, "..", "..", "data", "hsc_cosmos_35.npz") 

67 filename = os.path.abspath(filename) 

68 data = np.load(filename) 

69 self.data = data 

70 model_psf = integrated_circular_gaussian(sigma=0.8) 

71 self.detect = np.sum(data["images"], axis=0) 

72 self.centers = np.array([data["catalog"]["y"], data["catalog"]["x"]]).T + np.array(yx0) 

73 bands = data["filters"] 

74 self.observation = Observation( 

75 Image(data["images"], bands=bands, yx0=yx0), 

76 Image(data["variance"], bands=bands, yx0=yx0), 

77 Image(1 / data["variance"], bands=bands, yx0=yx0), 

78 data["psfs"], 

79 model_psf[None], 

80 bands=bands, 

81 ) 

82 

83 def tearDown(self): 

84 del self.data 

85 

86 def test_free_form_component(self): 

87 images = self.data["images"] 

88 

89 # Test with no thresholding (sparsity) 

90 sources = [] 

91 for i in range(5): 

92 component = FreeFormComponent( 

93 self.observation.bands, 

94 np.ones(5), 

95 images[i].copy(), 

96 self.observation.bbox, 

97 ) 

98 sources.append(Source([component])) 

99 

100 blend = Blend(sources, self.observation).fit_spectra() 

101 blend.parameterize(default_adaprox_parameterization) 

102 blend.fit(12, e_rel=1e-6) 

103 

104 # Test with thresholding (sparsity) 

105 sources = [] 

106 for i in range(5): 

107 component = FreeFormComponent( 

108 self.observation.bands, 

109 np.ones(5), 

110 images[i].copy(), 

111 self.observation.bbox, 

112 bg_rms=self.observation.noise_rms, 

113 bg_thresh=0.25, 

114 min_area=4, 

115 ) 

116 sources.append(Source([component])) 

117 

118 blend = Blend(sources, self.observation).fit_spectra() 

119 blend.parameterize(default_adaprox_parameterization) 

120 blend.fit(12, e_rel=1e-6) 

121 

122 # Test with peak centers specified 

123 sources = [] 

124 peaks = list(np.array([self.data["catalog"]["y"], self.data["catalog"]["x"]]).T.astype(int)) 

125 for i in range(5): 

126 component = FreeFormComponent( 

127 self.observation.bands, 

128 np.ones(5), 

129 images[i].copy(), 

130 self.observation.bbox, 

131 peaks=peaks, 

132 ) 

133 sources.append(Source([component])) 

134 

135 blend = Blend(sources, self.observation).fit_spectra() 

136 blend.parameterize(default_adaprox_parameterization) 

137 blend.fit(12, e_rel=1e-6) 

138 

139 # Tests for code blocks that are difficult to reach, 

140 # to complete test coverage 

141 component = blend.sources[-1].components[0] 

142 self.assertFalse(component.resize(self.observation.bbox)) 

143 component.morph[:] = 0 

144 component.prox_morph(component.morph) 

145 

146 

147class TestParametric(ScarletTestCase): 

148 def setUp(self) -> None: 

149 filename = os.path.join(__file__, "..", "..", "data", "hsc_cosmos_35.npz") 

150 filename = os.path.abspath(filename) 

151 data = np.load(filename) 

152 self.data = data 

153 self.model_psf = integrated_circular_gaussian(sigma=0.8) 

154 self.detect = np.sum(data["images"], axis=0) 

155 self.centers = np.array([data["catalog"]["y"], data["catalog"]["x"]]).T 

156 bands = data["filters"] 

157 self.observation = Observation( 

158 Image(data["images"], bands=bands), 

159 Image(data["variance"], bands=bands), 

160 Image(1 / data["variance"], bands=bands), 

161 data["psfs"], 

162 self.model_psf[None], 

163 bands=bands, 

164 ) 

165 

166 def tearDown(self): 

167 del self.data 

168 

169 def test_cartesian_frame(self): 

170 bbox = Box((31, 60), (1, 2)) 

171 frame = CartesianFrame(bbox) 

172 y = np.linspace(1, 31, 31) 

173 x = np.linspace(2, 61, 60) 

174 x, y = np.meshgrid(x, y) 

175 

176 self.assertBoxEqual(frame.bbox, bbox) 

177 self.assertTupleEqual(frame.shape, bbox.shape) 

178 assert_array_equal(frame.x_grid, x) 

179 assert_array_equal(frame.y_grid, y) 

180 self.assertIsNone(frame._r2) 

181 self.assertIsNone(frame._r) 

182 

183 def test_ellipse_frame(self): 

184 y0 = 23 

185 x0 = 36 

186 major = 3 

187 minor = 2 

188 theta = 2 * np.pi / 3 

189 bbox = Box((50, 39), (2, 7)) 

190 r_min = 1e-20 

191 frame = EllipseFrame( 

192 y0, 

193 x0, 

194 major, 

195 minor, 

196 theta, 

197 bbox, 

198 ) 

199 

200 self.assertEqual(frame.x0, x0) 

201 self.assertEqual(frame.y0, y0) 

202 self.assertEqual(frame.major, major) 

203 self.assertEqual(frame.minor, minor) 

204 self.assertEqual(frame.theta, theta) 

205 self.assertEqual(frame.bbox, bbox) 

206 

207 y = np.linspace(2, 51, 50) 

208 x = np.linspace(7, 45, 39) 

209 x, y = np.meshgrid(x, y) 

210 r2 = frame._xa**2 + frame._yb**2 + r_min**2 

211 r = np.sqrt(r2) 

212 

213 self.assertEqual(frame._sin, np.sin(theta)) 

214 self.assertEqual(frame._cos, np.cos(theta)) 

215 self.assertBoxEqual(frame.bbox, bbox) 

216 self.assertTupleEqual(frame.shape, bbox.shape) 

217 assert_array_equal(frame.x_grid, x) 

218 assert_array_equal(frame.y_grid, y) 

219 assert_array_equal(frame.r2_grid, r2) 

220 assert_array_equal(frame.r_grid, r) 

221 

222 # There doesn't seem to be much utility to testing the 

223 # gradients exactly, however we test that they all run 

224 # properly. 

225 gradients = ( 

226 frame.grad_x0, 

227 frame.grad_y0, 

228 frame.grad_major, 

229 frame.grad_minor, 

230 frame.grad_theta, 

231 ) 

232 

233 input_grad = np.ones(bbox.shape, dtype=float) 

234 original_grad = input_grad.copy() 

235 for grad in gradients: 

236 grad(input_grad, False) 

237 grad(input_grad, True) 

238 

239 # Make sure that none of the methods changed the input gradient 

240 assert_array_equal(input_grad, original_grad) 

241 

242 def test_parametric_component(self): 

243 observation = self.observation 

244 bands = observation.bands 

245 spectrum = np.ones((observation.n_bands,), dtype=float) 

246 frame = CartesianFrame(observation.bbox) 

247 

248 # Test integrated Gaussian PSF and sersic 

249 sources = [] 

250 for idx, center in enumerate(self.centers): 

251 # Get the integer center of the source 

252 cy, cx = int(np.round(center[0])), int(np.round(center[1])) 

253 # For now we use a fixed bounding box that is the size of 

254 # the observed PSF image 

255 bbox = Box((41, 41), origin=(cy - 20, cx - 20)) & observation.bbox 

256 # Keep track of the initial positions 

257 yi, xi = cy, cx 

258 # Restrict the values of the parameters 

259 _proxmin = np.array([yi - 2, xi - 2, 1e-1, 1e-1, -np.pi / 2, 1.1]) 

260 _proxmax = np.array([yi + 2, xi + 2, frame.shape[-2] / 2, frame.shape[-1] / 2, np.pi / 2, 3]) 

261 

262 __proxmin = np.array([yi - 2, xi - 2, 0.8]) 

263 __proxmax = np.array([yi + 2, xi + 2, 1.2]) 

264 

265 # Initialize a PSF-like component of the source using a 

266 # non-pixel integrated gaussian 

267 component = ParametricComponent( 

268 bands, 

269 bbox, 

270 spectrum=parameter(spectrum.copy()), 

271 morph_params=parameter(np.array([center[0], center[1], 0.8])), 

272 morph_func=models.integrated_gaussian, 

273 morph_grad=models.grad_integrated_gaussian, 

274 morph_prox=partial(models.bounded_prox, proxmin=__proxmin, proxmax=__proxmax), 

275 morph_step=np.array([1e-2, 1e-2, 1e-2]), 

276 prox_spectrum=lambda x: x, 

277 ) 

278 # Define the component to use ADAPROX as the optimizer 

279 components = [component] 

280 

281 # Initialize an n=1 sersic component 

282 component = EllipticalParametricComponent( 

283 bands, 

284 bbox, 

285 spectrum=parameter(spectrum.copy()), 

286 morph_params=parameter(np.array([center[0], center[1], 2 * 1.2**2, 2 * 1.2**2, 0.0, 1])), 

287 morph_func=models.sersic, 

288 morph_grad=models.grad_sersic, 

289 morph_prox=partial(models.bounded_prox, proxmin=_proxmin, proxmax=_proxmax), 

290 morph_step=np.array([1e-2, 1e-2, 1e-3, 1e-3, 1e-2, 1e-2]), 

291 ) 

292 # Define the component to use ADAPROX as the optimizer 

293 components.append(component) 

294 

295 component = EllipticalParametricComponent( 

296 bands, 

297 bbox, 

298 spectrum=parameter(spectrum.copy()), 

299 morph_params=parameter(np.array([center[0], center[1], 2 * 1.2**2, 2 * 1.2**2, 0.0, 1])), 

300 morph_func=models.sersic, 

301 morph_grad=models.grad_sersic, 

302 morph_prox=partial(models.bounded_prox, proxmin=_proxmin, proxmax=_proxmax), 

303 morph_step=np.array([1e-2, 1e-2, 1e-3, 1e-3, 1e-2, 1e-2]), 

304 ) 

305 # Define the component to use ADAPROX as the optimizer 

306 components.append(component) 

307 

308 # Create a new source using the two components 

309 sources.append(Source(components)) 

310 

311 # Fit the models 

312 blend = Blend(sources, observation) 

313 blend.parameterize(parameterize) 

314 blend.fit_spectra() 

315 

316 # Check properties of a component 

317 component = cast(ParametricComponent, blend.components[0]) 

318 self.assertTupleEqual(component.peak, tuple(self.centers[0])) 

319 self.assertEqual(component.y0, component._params.x[0]) 

320 self.assertEqual(component.x0, component._params.x[1]) 

321 assert_array_equal(component.morph_step, np.array([1e-2, 1e-2, 1e-2])) 

322 

323 component = cast(EllipticalParametricComponent, blend.components[1]) 

324 self.assertTupleEqual(component.peak, tuple(self.centers[0])) 

325 self.assertEqual(component.y0, component._params.x[0]) 

326 self.assertEqual(component.x0, component._params.x[1]) 

327 assert_array_equal(component.semi_major, component._params.x[2]) 

328 assert_array_equal(component.semi_minor, component._params.x[3]) 

329 assert_array_equal(component.theta, component._params.x[4]) 

330 assert_array_equal(component.ellipse_params, component._params.x[:5]) 

331 assert_array_equal(component.radial_params, component._params.x[5:]) 

332 self.assertIsNotNone(component.morph_prox) 

333 self.assertIsNotNone(component.morph_grad) 

334 

335 blend.fit(12) 

336 

337 # Test elliptical and circular Gaussian models 

338 sources = [] 

339 for idx, center in enumerate(self.centers): 

340 # Get the integer center of the source 

341 cy, cx = int(np.round(center[0])), int(np.round(center[1])) 

342 # For now we use a fixed bounding box that is the size of 

343 # the observed PSF image 

344 bbox = Box((41, 41), origin=(cy - 20, cx - 20)) & observation.bbox 

345 # Keep track of the initial positions 

346 yi, xi = cy, cx 

347 # Restrict the values of the parameters 

348 _proxmin = np.array([yi - 2, xi - 2, 1e-1, 1e-1, -np.pi / 2]) 

349 _proxmax = np.array([yi + 2, xi + 2, frame.shape[-2] / 2, frame.shape[-1] / 2, np.pi / 2]) 

350 

351 __proxmin = np.array([yi - 2, xi - 2]) 

352 __proxmax = np.array([yi + 2, xi + 2]) 

353 

354 # Initialize a PSF-like component of the source using a 

355 # non-pixel integrated gaussian 

356 component = ParametricComponent( 

357 bands, 

358 bbox, 

359 spectrum=parameter(spectrum.copy()), 

360 morph_params=parameter(np.array([center[0], center[1]])), 

361 morph_func=partial(models.circular_gaussian, sigma=0.8), 

362 morph_grad=partial(models.grad_circular_gaussian, sigma=0.8), 

363 morph_prox=partial(models.bounded_prox, proxmin=__proxmin, proxmax=__proxmax), 

364 morph_step=np.array([1e-2, 1e-2, 1e-2]), 

365 ) 

366 # Define the component to use ADAPROX as the optimizer 

367 components = [component] 

368 

369 # Initialize an n=1 sersic component 

370 component = EllipticalParametricComponent( 

371 bands, 

372 bbox, 

373 spectrum=parameter(spectrum.copy()), 

374 morph_params=parameter(np.array([center[0], center[1], 2 * 1.2**2, 2 * 1.2**2, 0.0])), 

375 morph_func=models.gaussian2d, 

376 morph_grad=models.grad_gaussian2, 

377 morph_prox=partial(models.bounded_prox, proxmin=_proxmin, proxmax=_proxmax), 

378 morph_step=np.array([1e-2, 1e-2, 1e-3, 1e-3, 1e-2, 1e-2]), 

379 ) 

380 # Define the component to use ADAPROX as the optimizer 

381 components.append(component) 

382 # Create a new source using the two components 

383 sources.append(Source(components)) 

384 

385 # Fit the models 

386 blend = Blend(sources, observation) 

387 blend.parameterize(parameterize) 

388 blend.fit_spectra() 

389 blend.fit(12) 

390 

391 def test_psf_fitting(self): 

392 # Use flat weights for FISTA optimization 

393 weights = np.ones(self.data["images"].shape) 

394 

395 monotonicity = Monotonicity((101, 101)) 

396 

397 observation = FittedPsfObservation( 

398 self.data["images"], 

399 self.data["variance"], 

400 weights, 

401 self.data["psfs"], 

402 self.model_psf[None], 

403 bands=self.data["filters"], 

404 ) 

405 

406 def obs_params(cls): 

407 if isinstance(cls, FittedPsfObservation): 

408 cls._fitted_kernel = FistaParameter(cls._fitted_kernel.x, step=1e-2) 

409 

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

411 blend = FittedPsfBlend(init.sources, observation).fit_spectra() 

412 blend.parameterize(default_fista_parameterization) 

413 cast(FittedPsfObservation, blend.observation).parameterize(obs_params) 

414 blend.fit(12, e_rel=1e-4)