Coverage for python/lsst/scarlet/lite/models/parametric.py: 33%

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 

24__all__ = [ 

25 "bounded_prox", 

26 "gaussian2d", 

27 "grad_gaussian2", 

28 "circular_gaussian", 

29 "grad_circular_gaussian", 

30 "integrated_gaussian", 

31 "grad_integrated_gaussian", 

32 "sersic", 

33 "grad_sersic", 

34 "CartesianFrame", 

35 "EllipseFrame", 

36 "ParametricComponent", 

37 "EllipticalParametricComponent", 

38] 

39 

40from typing import TYPE_CHECKING, Callable, Sequence, cast 

41 

42import numpy as np 

43from scipy.special import erf 

44from scipy.stats import gamma 

45 

46from ..bbox import Box 

47from ..component import Component 

48from ..image import Image 

49from ..parameters import Parameter, parameter 

50 

51if TYPE_CHECKING: 

52 from ..io import ScarletComponentBaseData 

53 

54# Some operations fail at the origin in radial coordinates, 

55# so we make use of a very small offset. 

56MIN_RADIUS = 1e-20 

57 

58# Useful constants 

59SQRT_PI_2 = np.sqrt(np.pi / 2) 

60 

61# Stored sersic constants 

62SERSIC_B1 = gamma.ppf(0.5, 2) 

63 

64 

65class CartesianFrame: 

66 """A grid of X and Y values contained in a bbox""" 

67 

68 def __init__(self, bbox: Box): 

69 """ 

70 Parameters 

71 ---------- 

72 bbox: Box 

73 The bounding box that contains this frame. 

74 """ 

75 # Store the new bounding box 

76 self._bbox = bbox 

77 # Get the range of x and y 

78 yi, xi = bbox.start[-2:] 

79 yf, xf = bbox.stop[-2:] 

80 height, width = bbox.shape[-2:] 

81 y = np.linspace(yi, yf - 1, height) 

82 x = np.linspace(xi, xf - 1, width) 

83 # Create the grid used to create the image of the frame 

84 self._x, self._y = np.meshgrid(x, y) 

85 self._r = None 

86 self._r2 = None 

87 

88 @property 

89 def shape(self) -> tuple[int, ...]: 

90 """Shape of the frame.""" 

91 return self._bbox.shape 

92 

93 @property 

94 def bbox(self) -> Box: 

95 """Bounding box containing the frame""" 

96 return self._bbox 

97 

98 @property 

99 def x_grid(self) -> np.ndarray: 

100 """The grid of x-values for the entire frame""" 

101 return self._x 

102 

103 @property 

104 def y_grid(self) -> np.ndarray: 

105 """The grid of y-values for the entire frame""" 

106 return self._y 

107 

108 

109class EllipseFrame(CartesianFrame): 

110 """Frame to scale the radius based on the parameters of an ellipse 

111 

112 This frame is used to calculate the coordinates of the 

113 radius and radius**2 from a given center location, 

114 based on the semi-major axis, semi-minor axis, and rotation angle. 

115 It is also used to calculate the gradient wrt either the 

116 radius**2 or radius for all of the model parameters. 

117 

118 Parameters 

119 ---------- 

120 y0: float 

121 The y-center of the ellipse. 

122 x0: float 

123 The x-center of the ellipse. 

124 major: float 

125 The length of the semi-major axis. 

126 minor: float 

127 The length of the semi-minor axis. 

128 theta: float 

129 The counter-clockwise rotation angle 

130 from the semi-major axis. 

131 bbox: Box 

132 The bounding box that contains the entire frame. 

133 r_min: float 

134 The minimum value of the radius. 

135 This is used to prevent divergences that occur 

136 when calculating the gradient at radius == 0. 

137 """ 

138 

139 def __init__( 

140 self, 

141 y0: float, 

142 x0: float, 

143 major: float, 

144 minor: float, 

145 theta: float, 

146 bbox: Box, 

147 r_min: float = 1e-20, 

148 ): 

149 super().__init__(bbox) 

150 # Set some useful parameters for derivations 

151 sin = np.sin(theta) 

152 cos = np.cos(theta) 

153 

154 # Rotate into the frame with xMajor as the x-axis 

155 # and xMinor as the y-axis 

156 self._xMajor = (self._x - x0) * cos + (self._y - y0) * sin 

157 self._xMinor = -(self._x - x0) * sin + (self._y - y0) * cos 

158 # The scaled major and minor axes 

159 self._xa = self._xMajor / major 

160 self._yb = self._xMinor / minor 

161 

162 # Store parameters useful for gradient calculation 

163 self._y0, self._x0 = y0, x0 

164 self._theta = theta 

165 self._major = major 

166 self._minor = minor 

167 self._sin, self._cos = sin, cos 

168 self._bbox = bbox 

169 self._rMin = r_min 

170 # Store the scaled radius**2 and radius 

171 self._radius2 = self._xa**2 + self._yb**2 

172 self._radius: np.ndarray | None = None 

173 

174 def grad_x0(self, input_grad: np.ndarray, use_r2: bool) -> float: 

175 """The gradient of either the radius or radius**2 wrt. the x-center 

176 

177 Parameters 

178 ---------- 

179 input_grad: 

180 Gradient of the likelihood wrt the component model 

181 use_r2: 

182 Whether to calculate the gradient of the radius**2 

183 (``useR2==True``) or the radius (``useR2==False``). 

184 

185 Returns 

186 ------- 

187 result: 

188 The gradient of the likelihood wrt x0. 

189 """ 

190 grad = -self._xa * self._cos / self._major + self._yb * self._sin / self._minor 

191 if use_r2: 

192 grad *= 2 

193 else: 

194 r = self.r_grid 

195 grad *= 1 / r 

196 return np.sum(grad * input_grad) # type: ignore 

197 

198 def grad_y0(self, input_grad: np.ndarray, use_r2: bool) -> float: 

199 """The gradient of either the radius or radius**2 wrt. the y-center 

200 

201 Parameters 

202 ---------- 

203 input_grad: 

204 Gradient of the likelihood wrt the component model 

205 use_r2: 

206 Whether to calculate the gradient of the radius**2 

207 (``useR2==True``) or the radius (``useR2==False``). 

208 

209 Returns 

210 ------- 

211 result: 

212 The gradient of the likelihood wrt y0. 

213 """ 

214 grad = -(self._xa * self._sin / self._major + self._yb * self._cos / self._minor) 

215 if use_r2: 

216 grad *= 2 

217 else: 

218 r = self.r_grid 

219 grad *= 1 / r 

220 return np.sum(grad * input_grad) # type: ignore 

221 

222 def grad_major(self, input_grad: np.ndarray, use_r2: bool) -> float: 

223 """The gradient of either the radius or radius**2 wrt. 

224 the semi-major axis 

225 

226 Parameters 

227 ---------- 

228 input_grad: 

229 Gradient of the likelihood wrt the component model 

230 use_r2: 

231 Whether to calculate the gradient of the radius**2 

232 (``use_r2==True``) or the radius (``use_r2==False``). 

233 

234 Returns 

235 ------- 

236 result: 

237 The gradient of the likelihood wrt the semi-major axis. 

238 """ 

239 grad = -2 / self._major * self._xa**2 

240 if use_r2: 

241 grad *= 2 

242 else: 

243 r = self.r_grid 

244 grad *= 1 / r 

245 return np.sum(grad * input_grad) # type: ignore 

246 

247 def grad_minor(self, input_grad: np.ndarray, use_r2: bool) -> float: 

248 """The gradient of either the radius or radius**2 wrt. 

249 the semi-minor axis 

250 

251 Parameters 

252 ---------- 

253 input_grad: 

254 Gradient of the likelihood wrt the component model 

255 use_r2: 

256 Whether to calculate the gradient of the radius**2 

257 (``useR2==True``) or the radius (``useR2==False``). 

258 

259 Returns 

260 ------- 

261 result: 

262 The gradient of the likelihood wrt the semi-minor axis. 

263 """ 

264 grad = -2 / self._minor * self._yb**2 

265 if use_r2: 

266 grad *= 2 

267 else: 

268 r = self.r_grid 

269 grad *= 1 / r 

270 return np.sum(grad * input_grad) # type: ignore 

271 

272 def grad_theta(self, input_grad: np.ndarray, use_r2: bool) -> float: 

273 """The gradient of either the radius or radius**2 wrt. 

274 the rotation angle 

275 

276 Parameters 

277 ---------- 

278 input_grad: 

279 Gradient of the likelihood wrt the component model 

280 use_r2: 

281 Whether to calculate the gradient of the radius**2 

282 (``useR2==True``) or the radius (``useR2==False``). 

283 

284 Returns 

285 ------- 

286 result: 

287 The gradient of the likelihood wrt the rotation angle. 

288 """ 

289 grad = self._xa * self._xMinor / self._major - self._yb * self._xMajor / self._minor 

290 if use_r2: 

291 grad *= 2 

292 else: 

293 r = self.r_grid 

294 grad *= 1 / r 

295 return np.sum(grad * input_grad) # type: ignore 

296 

297 @property 

298 def x0(self) -> float: 

299 """The x-center""" 

300 return self._x0 

301 

302 @property 

303 def y0(self) -> float: 

304 """The y-center""" 

305 return self._y0 

306 

307 @property 

308 def major(self) -> float: 

309 """The semi-major axis""" 

310 return self._major 

311 

312 @property 

313 def minor(self) -> float: 

314 """The semi-minor axis""" 

315 return self._minor 

316 

317 @property 

318 def theta(self) -> float: 

319 """The rotation angle""" 

320 return self._theta 

321 

322 @property 

323 def bbox(self) -> Box: 

324 """The bounding box to hold the model""" 

325 return self._bbox 

326 

327 @property 

328 def r_grid(self) -> np.ndarray: 

329 """The radial coordinates of each pixel""" 

330 if self._radius is None: 

331 self._radius = np.sqrt(self.r2_grid) 

332 return self._radius 

333 

334 @property 

335 def r2_grid(self) -> np.ndarray: 

336 """The radius squared located at each pixel""" 

337 return self._radius2 + self._rMin**2 

338 

339 

340def gaussian2d(params: np.ndarray, ellipse: EllipseFrame) -> np.ndarray: 

341 """Model of a 2D elliptical gaussian 

342 

343 Parameters 

344 ---------- 

345 params: 

346 The parameters of the function. 

347 In this case there are none outside of the ellipticity 

348 ellipse: 

349 The ellipse parameters to scale the radius in all directions. 

350 

351 Returns 

352 ------- 

353 result: 

354 The 2D guassian for the given ellipse parameters 

355 """ 

356 return np.exp(-ellipse.r2_grid) 

357 

358 

359def grad_gaussian2( 

360 input_grad: np.ndarray, 

361 params: np.ndarray, 

362 morph: np.ndarray, 

363 spectrum: np.ndarray, 

364 ellipse: EllipseFrame, 

365) -> np.ndarray: 

366 """Gradient of the the component model wrt the Gaussian 

367 morphology parameters 

368 

369 Parameters 

370 ---------- 

371 input_grad: 

372 Gradient of the likelihood wrt the component model 

373 params: 

374 The parameters of the morphology. 

375 morph: 

376 The model of the morphology. 

377 spectrum: 

378 The model of the spectrum. 

379 ellipse: 

380 The ellipse parameters to scale the radius in all directions. 

381 """ 

382 # Calculate the gradient of the likelihod 

383 # wrt the Gaussian e^-r**2 

384 _grad = -morph * np.einsum("i,i...", spectrum, input_grad) 

385 d_y0 = ellipse.grad_y0(_grad, True) 

386 d_x0 = ellipse.grad_x0(_grad, True) 

387 d_sigma_y = ellipse.grad_major(_grad, True) 

388 d_sigma_x = ellipse.grad_minor(_grad, True) 

389 d_theta = ellipse.grad_theta(_grad, True) 

390 return np.array([d_y0, d_x0, d_sigma_y, d_sigma_x, d_theta], dtype=params.dtype) 

391 

392 

393def circular_gaussian(center: Sequence[int], frame: CartesianFrame, sigma: float) -> np.ndarray: 

394 """Model of a circularly symmetric Gaussian 

395 

396 Parameters 

397 ---------- 

398 center: 

399 The center of the Gaussian. 

400 frame: 

401 The frame in which to generate the image of the circular Gaussian 

402 sigma: 

403 The standard deviation. 

404 

405 Returns 

406 ------- 

407 result: 

408 The image of the circular Gaussian. 

409 """ 

410 y0, x0 = center[:2] 

411 two_sigma = 2 * sigma 

412 r2 = ((frame.x_grid - x0) / two_sigma) ** 2 + ((frame.y_grid - y0) / two_sigma) ** 2 

413 return np.exp(-r2) 

414 

415 

416def grad_circular_gaussian( 

417 input_grad: np.ndarray, 

418 params: np.ndarray, 

419 morph: np.ndarray, 

420 spectrum: np.ndarray, 

421 frame: CartesianFrame, 

422 sigma: float, 

423) -> np.ndarray: 

424 """Gradient of the the component model wrt the Gaussian 

425 morphology parameters 

426 

427 Parameters 

428 ---------- 

429 input_grad: 

430 Gradient of the likelihood wrt the component model 

431 params: 

432 The parameters of the morphology. 

433 morph: 

434 The model of the morphology. 

435 spectrum: 

436 The model of the spectrum. 

437 frame: 

438 The frame in which to generate the image of the circular Gaussian. 

439 sigma: 

440 The standard deviation. 

441 """ 

442 # Calculate the gradient of the likelihod 

443 # wrt the Gaussian e^-r**2 

444 

445 _grad = -morph * np.einsum("i,i...", spectrum, input_grad) 

446 

447 y0, x0 = params[:2] 

448 d_y0 = -2 * np.sum((frame.y_grid - y0) * _grad) 

449 d_x0 = -2 * np.sum((frame.x_grid - x0) * _grad) 

450 return np.array([d_y0, d_x0], dtype=params.dtype) 

451 

452 

453def integrated_gaussian(params: np.ndarray, frame: CartesianFrame): 

454 """Model of a circularly symmetric Gaussian integrated over pixels 

455 

456 This differs from `circularGaussian` because the gaussian function 

457 is integrated over each pixel to replicate the pixelated image 

458 version of a Gaussian function. 

459 

460 Parameters 

461 ---------- 

462 params: 

463 The center of the Gaussian. 

464 frame: 

465 The frame in which to generate the image of the circular Gaussian 

466 

467 Returns 

468 ------- 

469 result: 

470 The image of the circular Gaussian. 

471 """ 

472 # Unpack the parameters and define constants 

473 y0, x0, sigma = params 

474 r = np.sqrt((frame.x_grid - x0) ** 2 + (frame.y_grid - y0) ** 2) 

475 sqrt_c = 1 / np.sqrt(2) / sigma 

476 # Integrate from half a pixel left and right 

477 lhs = erf((r - 0.5) * sqrt_c) 

478 rhs = erf((r + 0.5) * sqrt_c) 

479 z = 0.5 * np.sqrt(np.pi) / sqrt_c * (rhs - lhs) 

480 return z 

481 

482 

483def grad_integrated_gaussian( 

484 input_grad: np.ndarray, 

485 params: np.ndarray, 

486 morph: np.ndarray, 

487 spectrum: np.ndarray, 

488 frame: CartesianFrame, 

489) -> np.ndarray: 

490 """Gradient of the component model wrt the Gaussian 

491 morphology parameters 

492 

493 Parameters 

494 ---------- 

495 input_grad: 

496 Gradient of the likelihood wrt the component model parameters. 

497 params: 

498 The parameters of the morphology. 

499 morph: 

500 The model of the morphology. 

501 spectrum: 

502 The model of the spectrum. 

503 frame: 

504 The frame in which to generate the image of the circular Gaussian. 

505 

506 Returns 

507 ------- 

508 result: 

509 The gradient of the component morphology. 

510 """ 

511 # Calculate the gradient of the likelihood 

512 # wrt the Gaussian e^-r**2 

513 _grad = np.einsum("i,i...", spectrum, input_grad) 

514 

515 # Extract the parameters 

516 y0, x0, sigma = params 

517 # define useful constants 

518 x = frame.x_grid - x0 

519 y = frame.y_grid - y0 

520 c = 0.5 / sigma**2 

521 sqrt_c = np.sqrt(c) 

522 # Add a small constant to the radius to prevent a divergence at r==0 

523 r = np.sqrt(x**2 + y**2) + MIN_RADIUS 

524 # Shift half a pixel in each direction for the integration 

525 r1 = r - 0.5 

526 r2 = r + 0.5 

527 # Calculate the gradient of the ERF wrt. each shifted radius 

528 d_model1 = np.exp(-c * r1**2) 

529 d_model2 = np.exp(-c * r2**2) 

530 # Calculate the gradients of the parameters 

531 d_x0 = np.sum(-x / r * (d_model2 - d_model1) * _grad) 

532 d_y0 = np.sum(-y / r * (d_model2 - d_model1) * _grad) 

533 d_sigma1 = -(r1 * d_model1 / sigma - SQRT_PI_2 * erf(r1 * sqrt_c)) 

534 d_sigma2 = -(r2 * d_model2 / sigma - SQRT_PI_2 * erf(r2 * sqrt_c)) 

535 d_sigma = np.sum((d_sigma2 - d_sigma1) * _grad) 

536 

537 return np.array([d_y0, d_x0, d_sigma]) 

538 

539 

540def bounded_prox(params: np.ndarray, proxmin: np.ndarray, proxmax: np.ndarray) -> np.ndarray: 

541 """A bounded proximal operator 

542 

543 This function updates `params` in place. 

544 

545 Parameters 

546 ---------- 

547 params: 

548 The array of parameters to constrain. 

549 proxmin: 

550 The array of minimum values for each parameter. 

551 proxmax: 

552 The array of maximum values for each parameter. 

553 

554 Returns 

555 ------- 

556 result: 

557 The updated parameters. 

558 """ 

559 cuts = params < proxmin 

560 params[cuts] = proxmin[cuts] 

561 cuts = params > proxmax 

562 params[cuts] = proxmax[cuts] 

563 return params 

564 

565 

566def sersic(params: np.ndarray, ellipse: EllipseFrame) -> np.ndarray: 

567 """Generate a Sersic Model. 

568 

569 Parameters 

570 ---------- 

571 params: 

572 The parameters of the function. 

573 In this case the only parameter is the sersic index ``n``. 

574 ellipse: 

575 The ellipse parameters to scale the radius in all directions. 

576 

577 Returns 

578 ------- 

579 result: 

580 The model for the given ellipse parameters 

581 """ 

582 (n,) = params 

583 

584 r = ellipse.r_grid 

585 

586 if n == 1: 

587 result = np.exp(-SERSIC_B1 * r) 

588 else: 

589 bn = gamma.ppf(0.5, 2 * n) 

590 result = np.exp(-bn * (r ** (1 / n) - 1)) 

591 return result 

592 

593 

594def grad_sersic( 

595 input_grad: np.ndarray, 

596 params: np.ndarray, 

597 morph: np.ndarray, 

598 spectrum: np.ndarray, 

599 ellipse: EllipseFrame, 

600) -> np.ndarray: 

601 """Gradient of the component model wrt the morphology parameters 

602 

603 Parameters 

604 ---------- 

605 input_grad: 

606 Gradient of the likelihood wrt the component model 

607 params: 

608 The parameters of the morphology. 

609 morph: 

610 The model of the morphology. 

611 spectrum: 

612 The model of the spectrum. 

613 ellipse: 

614 The ellipse parameters to scale the radius in all directions. 

615 """ 

616 n = params[5] 

617 bn = gamma.ppf(0.5, 2 * n) 

618 if n == 1: 

619 # Use a simplified model for faster calculation 

620 d_exp = -SERSIC_B1 * morph 

621 else: 

622 r = ellipse.r_grid 

623 d_exp = -bn / n * morph * r ** (1 / n - 1) 

624 

625 _grad = np.einsum("i,i...", spectrum, input_grad) 

626 d_n = np.sum(_grad * bn * morph * ellipse.r_grid ** (1 / n) * np.log10(ellipse.r_grid) / n**2) 

627 _grad = _grad * d_exp 

628 d_y0 = ellipse.grad_y0(_grad, False) 

629 d_x0 = ellipse.grad_x0(_grad, False) 

630 d_sigma_y = ellipse.grad_major(_grad, False) 

631 d_sigma_x = ellipse.grad_minor(_grad, False) 

632 d_theta = ellipse.grad_theta(_grad, False) 

633 return np.array([d_y0, d_x0, d_sigma_y, d_sigma_x, d_theta, d_n], dtype=params.dtype) 

634 

635 

636class ParametricComponent(Component): 

637 """A parametric model of an astrophysical source 

638 

639 Parameters 

640 ---------- 

641 bands: 

642 The bands used in the model. 

643 bbox: 

644 The bounding box that holds the model. 

645 spectrum: 

646 The spectrum of the component. 

647 morph_params: 

648 The parameters of the morphology. 

649 morph_func: 

650 The function to generate the 2D morphology image 

651 based on `morphParams`. 

652 morph_grad: 

653 The function to calculate the gradient of the 

654 likelihood wrt the morphological parameters. 

655 morph_prox: 

656 The proximal operator for the morphology parameters. 

657 morph_step: 

658 The function that calculates the gradient of the 

659 morphological model. 

660 prox_spectrum: 

661 Proximal operator for the spectrum. 

662 If `prox_spectrum` is `None` then the default proximal 

663 operator `self.prox_spectrum` is used. 

664 floor: 

665 The minimum value of the spectrum, used to prevent 

666 divergences in the gradients. 

667 """ 

668 

669 def __init__( 

670 self, 

671 bands: tuple, 

672 bbox: Box, 

673 spectrum: Parameter | np.ndarray, 

674 morph_params: Parameter | np.ndarray, 

675 morph_func: Callable, 

676 morph_grad: Callable, 

677 morph_prox: Callable, 

678 morph_step: Callable | np.ndarray, 

679 prox_spectrum: Callable | None = None, 

680 floor: float = 1e-20, 

681 ): 

682 super().__init__(bands=bands, bbox=bbox) 

683 

684 self._spectrum = parameter(spectrum) 

685 self._params = parameter(morph_params) 

686 self._func = morph_func 

687 self._morph_grad = morph_grad 

688 self._morph_prox = morph_prox 

689 self._morph_step = morph_step 

690 self._bbox = bbox 

691 if prox_spectrum is None: 

692 self._prox_spectrum: Callable = self.prox_spectrum 

693 else: 

694 self._prox_spectrum = prox_spectrum 

695 self.floor = floor 

696 

697 @property 

698 def peak(self) -> tuple[float, float]: 

699 """The center of the component""" 

700 return self.y0, self.x0 

701 

702 @property 

703 def y0(self) -> float: 

704 """The y-center of the component""" 

705 return self._params.x[0] 

706 

707 @property 

708 def x0(self) -> float: 

709 """The x-center of the component""" 

710 return self._params.x[1] 

711 

712 @property 

713 def spectrum(self) -> np.ndarray: 

714 """The array of spectrum values""" 

715 return self._spectrum.x 

716 

717 @property 

718 def frame(self) -> CartesianFrame: 

719 """The coordinate system that contains the model""" 

720 return CartesianFrame(self._bbox) 

721 

722 @property 

723 def radial_params(self) -> np.ndarray: 

724 """The parameters used to model the radial function""" 

725 return self._params.x 

726 

727 def _get_morph(self, frame: CartesianFrame | None = None) -> np.ndarray: 

728 """The 2D image of the morphology 

729 

730 This callable generates an image of the morphology 

731 in the given frame. 

732 

733 Parameters 

734 ---------- 

735 frame: 

736 The frame (bounding box, pixel grid) that the image is 

737 placed in. 

738 

739 Returns 

740 ------- 

741 result: 

742 The image of the morphology in the `frame`. 

743 """ 

744 if frame is None: 

745 frame = self.frame 

746 return self._func(self.radial_params, frame) 

747 

748 @property 

749 def morph(self, frame: CartesianFrame | None = None) -> np.ndarray: 

750 """The morphological model""" 

751 return self._get_morph(frame) 

752 

753 @property 

754 def prox_morph(self) -> Callable: 

755 """The function used to constrain the morphological model""" 

756 return self._morph_prox 

757 

758 @property 

759 def grad_morph(self) -> Callable: 

760 """The function that calculates the gradient of the 

761 morphological model 

762 """ 

763 return self._morph_grad 

764 

765 @property 

766 def morph_step(self) -> Callable: 

767 """The function that calculates the gradient of the 

768 morphological model 

769 """ 

770 return cast(Callable, self._morph_step) 

771 

772 def get_model(self, frame: CartesianFrame | None = None) -> Image: 

773 """Generate the full model for this component""" 

774 model = self.spectrum[:, None, None] * self._get_morph(frame)[None, :, :] 

775 return Image(model, bands=self.bands, yx0=cast(tuple[int, int], self.bbox.origin[-2:])) 

776 

777 def prox_spectrum(self, spectrum: np.ndarray) -> np.ndarray: 

778 """Apply a prox-like update to the spectrum 

779 

780 Parameters 

781 ---------- 

782 spectrum: 

783 The spectrum of the model. 

784 """ 

785 # prevent divergent spectrum 

786 spectrum[spectrum < self.floor] = self.floor 

787 return spectrum 

788 

789 def grad_spectrum(self, input_grad: np.ndarray, spectrum: np.ndarray, morph: np.ndarray) -> np.ndarray: 

790 """Gradient of the spectrum wrt. the component model 

791 

792 Parameters 

793 ---------- 

794 input_grad: 

795 Gradient of the likelihood wrt the component model 

796 spectrum: 

797 The model of the spectrum. 

798 morph: 

799 The model of the morphology. 

800 

801 Returns 

802 ------- 

803 result: 

804 The gradient of the likelihood wrt. the spectrum. 

805 """ 

806 return np.einsum("...jk,jk", input_grad, morph) 

807 

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

809 """Update the component parameters from an input gradient 

810 

811 Parameters 

812 ---------- 

813 it: 

814 The current iteration of the optimizer. 

815 input_grad: 

816 Gradient of the likelihood wrt the component model 

817 """ 

818 spectrum = self.spectrum.copy() 

819 morph = self.morph 

820 self._spectrum.update(it, input_grad, morph) 

821 self._params.update(it, input_grad, morph, spectrum, self.frame) 

822 

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

824 """Resize the box that contains the model 

825 

826 Not yet implemented, so for now the model box 

827 does not grow. If this is ever implemented in production, 

828 in the long run this will be based on a cutoff value for the model. 

829 """ 

830 return False 

831 

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

833 """Convert the component parameter arrays into Parameter instances""" 

834 # Update the spectrum and morph in place 

835 parameterization(self) 

836 # update the parameters 

837 self._spectrum.grad = self.grad_spectrum 

838 self._spectrum.prox = self.prox_spectrum 

839 self._params.grad = self.grad_morph 

840 self._params.prox = self.prox_morph 

841 

842 def to_data(self) -> ScarletComponentBaseData: 

843 raise NotImplementedError("Saving elliptical parametric components is not yet implemented") 

844 

845 

846class EllipticalParametricComponent(ParametricComponent): 

847 """A radial density/surface brightness profile with elliptical symmetry 

848 

849 Parameters 

850 ---------- 

851 bands: 

852 The bands used in the model. 

853 bbox: 

854 The bounding box that holds this component model. 

855 spectrum: 

856 The spectrum of the component. 

857 morph_params: 

858 The parameters passed to `morph_func` to 

859 generate the morphology in image space. 

860 morph_func: 

861 The function to generate the morphology 

862 based on `morphParams`. 

863 morph_grad: 

864 The function to calculate the gradient of the 

865 likelihood wrt the morphological parameters. 

866 morph_prox: 

867 The proximal operator for the morphology parameters. 

868 prox_spectrum: 

869 Proximal operator for the spectrum. 

870 If `prox_spectrum` is `None` then the default proximal 

871 operator `self.prox_spectrum` is used. 

872 floor: 

873 The minimum value of the spectrum, used to prevent 

874 divergences in the gradients. 

875 """ 

876 

877 def __init__( 

878 self, 

879 bands: tuple, 

880 bbox: Box, 

881 spectrum: Parameter | np.ndarray, 

882 morph_params: Parameter | np.ndarray, 

883 morph_func: Callable, 

884 morph_grad: Callable, 

885 morph_prox: Callable, 

886 morph_step: Callable | np.ndarray, 

887 prox_spectrum: Callable | None = None, 

888 floor: float = 1e-20, 

889 ): 

890 super().__init__( 

891 bands=bands, 

892 bbox=bbox, 

893 spectrum=spectrum, 

894 morph_params=morph_params, 

895 morph_func=morph_func, 

896 morph_grad=morph_grad, 

897 morph_prox=morph_prox, 

898 morph_step=morph_step, 

899 prox_spectrum=prox_spectrum, 

900 floor=floor, 

901 ) 

902 

903 @property 

904 def semi_major(self) -> float: 

905 """The length of the semi-major axis of the model""" 

906 return self._params.x[2] 

907 

908 @property 

909 def semi_minor(self) -> float: 

910 """The length of the semi-minor axis of the model""" 

911 return self._params.x[3] 

912 

913 @property 

914 def theta(self) -> float: 

915 """The counter-clockwise rotation angle of the model from the 

916 x-axis. 

917 """ 

918 return self._params.x[4] 

919 

920 @property 

921 def ellipse_params(self) -> np.ndarray: 

922 """The parameters used to generate the scaled radius""" 

923 return self._params.x[:5] 

924 

925 @property 

926 def radial_params(self) -> np.ndarray: 

927 """The parameters used to model the radial function""" 

928 return self._params.x[5:] 

929 

930 @property 

931 def frame(self) -> EllipseFrame: