Coverage for tests/test_models.py: 10%

294 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-05-30 08:25 +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 Component, FactorizedComponent, default_adaprox_parameterization 

30from lsst.scarlet.lite.initialization import FactorizedInitialization 

31from lsst.scarlet.lite.models import ( 

32 CartesianFrame, 

33 EllipseFrame, 

34 EllipticalParametricComponent, 

35 FactorizedFreeFormComponent, 

36 FittedPsfBlend, 

37 FittedPsfObservation, 

38 ParametricComponent, 

39) 

40from lsst.scarlet.lite.operators import Monotonicity 

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

42from lsst.scarlet.lite.utils import integrated_circular_gaussian 

43from numpy.testing import assert_array_equal 

44from scipy.stats import gamma as gamma_dist 

45from utils import ScarletTestCase 

46 

47 

48def parameterize(component: Component): 

49 assert isinstance(component, ParametricComponent) 

50 component._spectrum = AdaproxParameter( 

51 component.spectrum, 

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

53 ) 

54 component._params = AdaproxParameter( 

55 component._params.x, 

56 step=1e-2, 

57 ) 

58 

59 

60class TestFreeForm(ScarletTestCase): 

61 def setUp(self) -> None: 

62 yx0 = (1000, 2000) 

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

64 filename = os.path.abspath(filename) 

65 data = np.load(filename) 

66 self.data = data 

67 model_psf = integrated_circular_gaussian(sigma=0.8) 

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

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

70 bands = data["filters"] 

71 self.observation = Observation( 

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

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

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

75 data["psfs"], 

76 model_psf[None], 

77 bands=bands, 

78 ) 

79 

80 def tearDown(self): 

81 del self.data 

82 

83 def test_free_form_component(self): 

84 images = self.data["images"] 

85 

86 # Test with no thresholding (sparsity) 

87 sources = [] 

88 for i in range(5): 

89 component = FactorizedFreeFormComponent( 

90 self.observation.bands, 

91 np.ones(5), 

92 images[i].copy(), 

93 self.observation.bbox, 

94 ) 

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

96 

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

98 blend.parameterize(default_adaprox_parameterization) 

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

100 

101 # Test with thresholding (sparsity) 

102 sources = [] 

103 for i in range(5): 

104 component = FactorizedFreeFormComponent( 

105 self.observation.bands, 

106 np.ones(5), 

107 images[i].copy(), 

108 self.observation.bbox, 

109 bg_rms=self.observation.noise_rms, 

110 bg_thresh=0.25, 

111 min_area=4, 

112 ) 

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

114 

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

116 blend.parameterize(default_adaprox_parameterization) 

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

118 

119 # Test with peak centers specified 

120 sources = [] 

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

122 for i in range(5): 

123 component = FactorizedFreeFormComponent( 

124 self.observation.bands, 

125 np.ones(5), 

126 images[i].copy(), 

127 self.observation.bbox, 

128 peaks=peaks, 

129 ) 

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

131 

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

133 blend.parameterize(default_adaprox_parameterization) 

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

135 

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

137 # to complete test coverage 

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

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

140 component.morph[:] = 0 

141 component.prox_morph(component.morph) 

142 

143 

144class TestParametric(ScarletTestCase): 

145 def setUp(self) -> None: 

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

147 filename = os.path.abspath(filename) 

148 data = np.load(filename) 

149 self.data = data 

150 self.model_psf = integrated_circular_gaussian(sigma=0.8) 

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

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

153 bands = data["filters"] 

154 self.observation = Observation( 

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

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

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

158 data["psfs"], 

159 self.model_psf[None], 

160 bands=bands, 

161 ) 

162 

163 def tearDown(self): 

164 del self.data 

165 

166 def test_cartesian_frame(self): 

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

168 frame = CartesianFrame(bbox) 

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

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

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

172 

173 self.assertBoxEqual(frame.bbox, bbox) 

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

175 assert_array_equal(frame.x_grid, x) 

176 assert_array_equal(frame.y_grid, y) 

177 self.assertIsNone(frame._r2) 

178 self.assertIsNone(frame._r) 

179 

180 def test_ellipse_frame(self): 

181 y0 = 23 

182 x0 = 36 

183 major = 3 

184 minor = 2 

185 theta = 2 * np.pi / 3 

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

187 r_min = 1e-20 

188 frame = EllipseFrame( 

189 y0, 

190 x0, 

191 major, 

192 minor, 

193 theta, 

194 bbox, 

195 ) 

196 

197 self.assertEqual(frame.x0, x0) 

198 self.assertEqual(frame.y0, y0) 

199 self.assertEqual(frame.major, major) 

200 self.assertEqual(frame.minor, minor) 

201 self.assertEqual(frame.theta, theta) 

202 self.assertEqual(frame.bbox, bbox) 

203 

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

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

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

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

208 r = np.sqrt(r2) 

209 

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

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

212 self.assertBoxEqual(frame.bbox, bbox) 

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

214 assert_array_equal(frame.x_grid, x) 

215 assert_array_equal(frame.y_grid, y) 

216 assert_array_equal(frame.r2_grid, r2) 

217 assert_array_equal(frame.r_grid, r) 

218 

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

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

221 # properly. 

222 gradients = ( 

223 frame.grad_x0, 

224 frame.grad_y0, 

225 frame.grad_major, 

226 frame.grad_minor, 

227 frame.grad_theta, 

228 ) 

229 

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

231 original_grad = input_grad.copy() 

232 for grad in gradients: 

233 grad(input_grad, False) 

234 grad(input_grad, True) 

235 

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

237 assert_array_equal(input_grad, original_grad) 

238 

239 def test_grad_sersic_n_index(self): 

240 """Finite-difference check of the Sersic gradient w.r.t. n. 

241 

242 Protects against regressions in the analytical gradient (e.g. the 

243 np.log10 vs np.log mistake fixed for audit finding C-1). The 

244 analytical gradient holds bn(n) fixed -- differentiating the 

245 inverse incomplete gamma function is intentionally omitted -- so 

246 the finite-difference comparison does the same. 

247 """ 

248 bbox = Box((33, 33), origin=(0, 0)) 

249 y0, x0 = 16.0, 16.0 

250 sigma_y, sigma_x = 5.0, 4.0 

251 theta = 0.3 

252 spectrum = np.array([1.0]) 

253 

254 def sersic_fixed_bn(n_value, bn_value, ellipse): 

255 r = ellipse.r_grid 

256 return np.exp(-bn_value * (r ** (1 / n_value) - 1)) 

257 

258 for n in (0.5, 1.5, 2.0, 3.0, 4.0): 

259 ellipse = EllipseFrame(y0, x0, sigma_y, sigma_x, theta, bbox) 

260 bn = gamma_dist.ppf(0.5, 2 * n) 

261 morph = sersic_fixed_bn(n, bn, ellipse) 

262 input_grad = np.ones((1,) + morph.shape) 

263 params = np.array([y0, x0, sigma_y, sigma_x, theta, n]) 

264 

265 d_n_analytical = models.grad_sersic(input_grad, params, morph, spectrum, ellipse)[5] 

266 

267 # Centered finite difference: df/dn ~ (f(n+e) - f(n-e)) / (2e) 

268 # where f(n) = sum(sersic_fixed_bn(n)). bn is held fixed in both 

269 # evaluations to match the analytical formulation, which omits 

270 # the dbn/dn contribution. 

271 eps = 1e-5 * max(1.0, abs(n)) 

272 loss_plus = np.sum(sersic_fixed_bn(n + eps, bn, ellipse)) 

273 loss_minus = np.sum(sersic_fixed_bn(n - eps, bn, ellipse)) 

274 d_n_fd = (loss_plus - loss_minus) / (2 * eps) 

275 

276 np.testing.assert_allclose( 

277 d_n_analytical, 

278 d_n_fd, 

279 rtol=1e-4, 

280 err_msg=f"grad_sersic d/dn mismatch at n={n}", 

281 ) 

282 

283 def test_grad_circular_gaussian(self): 

284 """Finite-difference check of the circular Gaussian gradient. 

285 

286 Protects against regressions in grad_circular_gaussian (audit 

287 finding C-2). The forward model is morph = exp(-r2) where 

288 r2 = ((x-x0)/(2*sigma))**2 + ((y-y0)/(2*sigma))**2, so the gradient 

289 must scale with 1/(2*sigma**2). A previous version of the code used 

290 a fixed factor of 2, ignoring sigma entirely. 

291 """ 

292 bbox = Box((33, 33), origin=(0, 0)) 

293 y0, x0 = 16.3, 16.7 

294 spectrum = np.array([1.0]) 

295 frame = CartesianFrame(bbox) 

296 

297 for sigma in (0.6, 0.8, 1.5, 3.0): 

298 morph = models.circular_gaussian((y0, x0), frame, sigma=sigma) 

299 input_grad = np.ones((1,) + morph.shape) 

300 params = np.array([y0, x0]) 

301 

302 d_analytical = models.grad_circular_gaussian( 

303 input_grad, params, morph, spectrum, frame, sigma=sigma 

304 ) 

305 

306 # Centered finite difference on f(y0, x0) = sum(morph(y0, x0)). 

307 eps = 1e-4 

308 f_yp = np.sum(models.circular_gaussian((y0 + eps, x0), frame, sigma=sigma)) 

309 f_ym = np.sum(models.circular_gaussian((y0 - eps, x0), frame, sigma=sigma)) 

310 f_xp = np.sum(models.circular_gaussian((y0, x0 + eps), frame, sigma=sigma)) 

311 f_xm = np.sum(models.circular_gaussian((y0, x0 - eps), frame, sigma=sigma)) 

312 d_fd = np.array([(f_yp - f_ym) / (2 * eps), (f_xp - f_xm) / (2 * eps)]) 

313 

314 np.testing.assert_allclose( 

315 d_analytical, 

316 d_fd, 

317 rtol=1e-4, 

318 atol=1e-8, 

319 err_msg=f"grad_circular_gaussian mismatch at sigma={sigma}", 

320 ) 

321 

322 def test_grad_gaussian2(self): 

323 """Finite-difference check of the elliptical 2D Gaussian 

324 gradient. 

325 

326 ``grad_gaussian2`` returns d/d{y0, x0, sigma_y, sigma_x, 

327 theta} of ``sum(spectrum * input_grad * morph)``. With 

328 ``spectrum = [1]`` and ``input_grad = ones``, the loss is 

329 just ``sum(morph)``. 

330 

331 Also exercises ``EllipseFrame.grad_major`` and 

332 ``grad_minor``: the ``-2/major * _xa**2`` base in those 

333 methods is already ``d(r**2)/d(major)``, while the 

334 ``grad_x0/y0/theta`` base is ``(1/2) * d(r**2)/d(...)``. 

335 Multiplying by 2 (``use_r2=True``) or by ``1/r`` 

336 (``use_r2=False``) was therefore producing 2x the correct 

337 gradient for the size parameters. 

338 """ 

339 bbox = Box((33, 33), origin=(0, 0)) 

340 spectrum = np.array([1.0]) 

341 # y0, x0, sigma_y, sigma_x, theta 

342 y0, x0, sigma_y, sigma_x, theta = 16.3, 16.7, 5.0, 4.0, 0.3 

343 empty_params = np.array([]) 

344 

345 ellipse = EllipseFrame(y0, x0, sigma_y, sigma_x, theta, bbox) 

346 morph = models.gaussian2d(empty_params, ellipse) 

347 input_grad = np.ones((1,) + morph.shape) 

348 params = np.array([y0, x0, sigma_y, sigma_x, theta]) 

349 d_analytical = models.grad_gaussian2(input_grad, params, morph, spectrum, ellipse) 

350 

351 # Smaller eps for theta because the morph is more sensitive 

352 # near theta=0.3 (radians), and a too-large eps flattens the 

353 # finite-difference signal. 

354 eps_list = [1e-4, 1e-4, 1e-4, 1e-4, 1e-5] 

355 d_fd = np.zeros(5) 

356 for i, eps in enumerate(eps_list): 

357 perturb = np.zeros(5) 

358 perturb[i] = eps 

359 yp, xp, syp, sxp, tp = params + perturb 

360 ym, xm, sym, sxm, tm = params - perturb 

361 f_plus = np.sum(models.gaussian2d(empty_params, EllipseFrame(yp, xp, syp, sxp, tp, bbox))) 

362 f_minus = np.sum(models.gaussian2d(empty_params, EllipseFrame(ym, xm, sym, sxm, tm, bbox))) 

363 d_fd[i] = (f_plus - f_minus) / (2 * eps) 

364 

365 np.testing.assert_allclose(d_analytical, d_fd, rtol=1e-3, atol=1e-7) 

366 

367 def test_grad_integrated_gaussian(self): 

368 """Finite-difference check of the integrated Gaussian 

369 gradient. ``grad_integrated_gaussian`` returns 

370 d/d{y0, x0, sigma}. 

371 """ 

372 bbox = Box((33, 33), origin=(0, 0)) 

373 spectrum = np.array([1.0]) 

374 frame = CartesianFrame(bbox) 

375 params = np.array([16.3, 16.7, 1.5]) 

376 morph = models.integrated_gaussian(params, frame) 

377 input_grad = np.ones((1,) + morph.shape) 

378 

379 d_analytical = models.grad_integrated_gaussian(input_grad, params, morph, spectrum, frame) 

380 

381 eps = 1e-4 

382 d_fd = np.zeros(3) 

383 for i in range(3): 

384 perturb = np.zeros(3) 

385 perturb[i] = eps 

386 f_plus = np.sum(models.integrated_gaussian(params + perturb, frame)) 

387 f_minus = np.sum(models.integrated_gaussian(params - perturb, frame)) 

388 d_fd[i] = (f_plus - f_minus) / (2 * eps) 

389 

390 np.testing.assert_allclose(d_analytical, d_fd, rtol=1e-3, atol=1e-7) 

391 

392 def test_grad_sersic_ellipse_params(self): 

393 """Finite-difference check of grad_sersic for the ellipse 

394 parameters. Index 5 (d/dn) is covered by 

395 ``test_grad_sersic_n_index``; this covers indices 0-4: 

396 d/d{y0, x0, sigma_y, sigma_x, theta}. ``bn(n)`` does not 

397 depend on these parameters, so the FD evaluation can use 

398 ``models.sersic`` directly. 

399 """ 

400 bbox = Box((33, 33), origin=(0, 0)) 

401 spectrum = np.array([1.0]) 

402 n = 1.5 

403 n_params = np.array([n]) 

404 y0, x0, sigma_y, sigma_x, theta = 16.0, 16.0, 5.0, 4.0, 0.3 

405 

406 ellipse = EllipseFrame(y0, x0, sigma_y, sigma_x, theta, bbox) 

407 morph = models.sersic(n_params, ellipse) 

408 input_grad = np.ones((1,) + morph.shape) 

409 params = np.array([y0, x0, sigma_y, sigma_x, theta, n]) 

410 d_analytical = models.grad_sersic(input_grad, params, morph, spectrum, ellipse)[:5] 

411 

412 eps_list = [1e-4, 1e-4, 1e-4, 1e-4, 1e-5] 

413 d_fd = np.zeros(5) 

414 ellipse_params = np.array([y0, x0, sigma_y, sigma_x, theta]) 

415 for i, eps in enumerate(eps_list): 

416 perturb = np.zeros(5) 

417 perturb[i] = eps 

418 yp, xp, syp, sxp, tp = ellipse_params + perturb 

419 ym, xm, sym, sxm, tm = ellipse_params - perturb 

420 f_plus = np.sum(models.sersic(n_params, EllipseFrame(yp, xp, syp, sxp, tp, bbox))) 

421 f_minus = np.sum(models.sersic(n_params, EllipseFrame(ym, xm, sym, sxm, tm, bbox))) 

422 d_fd[i] = (f_plus - f_minus) / (2 * eps) 

423 

424 np.testing.assert_allclose(d_analytical, d_fd, rtol=1e-3, atol=1e-7) 

425 

426 def test_parametric_component(self): 

427 observation = self.observation 

428 bands = observation.bands 

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

430 frame = CartesianFrame(observation.bbox) 

431 

432 # Test integrated Gaussian PSF and sersic 

433 sources = [] 

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

435 # Get the integer center of the source 

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

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

438 # the observed PSF image 

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

440 # Keep track of the initial positions 

441 yi, xi = cy, cx 

442 # Restrict the values of the parameters 

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

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

445 

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

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

448 

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

450 # non-pixel integrated gaussian 

451 component = ParametricComponent( 

452 bands, 

453 bbox, 

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

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

456 morph_func=models.integrated_gaussian, 

457 morph_grad=models.grad_integrated_gaussian, 

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

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

460 prox_spectrum=lambda x: x, 

461 ) 

462 # Define the component to use ADAPROX as the optimizer 

463 components = [component] 

464 

465 # Initialize an n=1 sersic component 

466 component = EllipticalParametricComponent( 

467 bands, 

468 bbox, 

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

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

471 morph_func=models.sersic, 

472 morph_grad=models.grad_sersic, 

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

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

475 ) 

476 # Define the component to use ADAPROX as the optimizer 

477 components.append(component) 

478 

479 component = EllipticalParametricComponent( 

480 bands, 

481 bbox, 

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

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

484 morph_func=models.sersic, 

485 morph_grad=models.grad_sersic, 

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

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

488 ) 

489 # Define the component to use ADAPROX as the optimizer 

490 components.append(component) 

491 

492 # Create a new source using the two components 

493 sources.append(Source(components)) 

494 

495 # Fit the models 

496 blend = Blend(sources, observation) 

497 blend.parameterize(parameterize) 

498 blend.fit_spectra() 

499 

500 # Check properties of a component 

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

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

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

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

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

506 

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

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

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

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

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

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

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

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

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

516 self.assertIsNotNone(component.morph_prox) 

517 self.assertIsNotNone(component.morph_grad) 

518 

519 blend.fit(12) 

520 

521 # Test elliptical and circular Gaussian models 

522 sources = [] 

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

524 # Get the integer center of the source 

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

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

527 # the observed PSF image 

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

529 # Keep track of the initial positions 

530 yi, xi = cy, cx 

531 # Restrict the values of the parameters 

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

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

534 

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

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

537 

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

539 # non-pixel integrated gaussian 

540 component = ParametricComponent( 

541 bands, 

542 bbox, 

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

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

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

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

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

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

549 ) 

550 # Define the component to use ADAPROX as the optimizer 

551 components = [component] 

552 

553 # Initialize an n=1 sersic component 

554 component = EllipticalParametricComponent( 

555 bands, 

556 bbox, 

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

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

559 morph_func=models.gaussian2d, 

560 morph_grad=models.grad_gaussian2, 

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

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

563 ) 

564 # Define the component to use ADAPROX as the optimizer 

565 components.append(component) 

566 # Create a new source using the two components 

567 sources.append(Source(components)) 

568 

569 # Fit the models 

570 blend = Blend(sources, observation) 

571 blend.parameterize(parameterize) 

572 blend.fit_spectra() 

573 blend.fit(12) 

574 

575 def test_psf_fitting(self): 

576 # Use flat weights for FISTA optimization 

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

578 

579 monotonicity = Monotonicity((101, 101)) 

580 

581 observation = FittedPsfObservation( 

582 self.data["images"], 

583 self.data["variance"], 

584 weights, 

585 self.data["psfs"], 

586 self.model_psf[None], 

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

588 ) 

589 

590 def fista_parameterization(component: Component): 

591 if isinstance(component, FactorizedComponent): 

592 component._spectrum = FistaParameter(component.spectrum, step=0.5) 

593 component._morph = FistaParameter(component.morph, step=0.5) 

594 else: 

595 if isinstance(component, FittedPsfObservation): 

596 component._fitted_kernel = FistaParameter(component._fitted_kernel.x, step=1e-2) 

597 

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

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

600 blend.parameterize(fista_parameterization) 

601 assert isinstance(cast(FittedPsfObservation, blend.observation)._fitted_kernel, FistaParameter) 

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