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

293 statements  

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

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 

22__all__ = [ 

23 "bounded_prox", 

24 "gaussian2d", 

25 "grad_gaussian2", 

26 "circular_gaussian", 

27 "grad_circular_gaussian", 

28 "integrated_gaussian", 

29 "grad_integrated_gaussian", 

30 "sersic", 

31 "grad_sersic", 

32 "CartesianFrame", 

33 "EllipseFrame", 

34 "ParametricComponent", 

35 "EllipticalParametricComponent", 

36] 

37 

38from typing import Callable, Sequence, cast 

39 

40import numpy as np 

41from scipy.special import erf 

42from scipy.stats import gamma 

43 

44from ..bbox import Box 

45from ..component import Component 

46from ..image import Image 

47from ..parameters import Parameter, parameter 

48 

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

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

51MIN_RADIUS = 1e-20 

52 

53# Useful constants 

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

55 

56# Stored sersic constants 

57SERSIC_B1 = gamma.ppf(0.5, 2) 

58 

59 

60class CartesianFrame: 

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

62 

63 def __init__(self, bbox: Box): 

64 """ 

65 Parameters 

66 ---------- 

67 bbox: Box 

68 The bounding box that contains this frame. 

69 """ 

70 # Store the new bounding box 

71 self._bbox = bbox 

72 # Get the range of x and y 

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

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

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

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

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

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

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

80 self._r = None 

81 self._r2 = None 

82 

83 @property 

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

85 """Shape of the frame.""" 

86 return self._bbox.shape 

87 

88 @property 

89 def bbox(self) -> Box: 

90 """Bounding box containing the frame""" 

91 return self._bbox 

92 

93 @property 

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

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

96 return self._x 

97 

98 @property 

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

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

101 return self._y 

102 

103 

104class EllipseFrame(CartesianFrame): 

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

106 

107 This frame is used to calculate the coordinates of the 

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

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

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

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

112 

113 Parameters 

114 ---------- 

115 y0: float 

116 The y-center of the ellipse. 

117 x0: float 

118 The x-center of the ellipse. 

119 major: float 

120 The length of the semi-major axis. 

121 minor: float 

122 The length of the semi-minor axis. 

123 theta: float 

124 The counter-clockwise rotation angle 

125 from the semi-major axis. 

126 bbox: Box 

127 The bounding box that contains the entire frame. 

128 r_min: float 

129 The minimum value of the radius. 

130 This is used to prevent divergences that occur 

131 when calculating the gradient at radius == 0. 

132 """ 

133 

134 def __init__( 

135 self, 

136 y0: float, 

137 x0: float, 

138 major: float, 

139 minor: float, 

140 theta: float, 

141 bbox: Box, 

142 r_min: float = 1e-20, 

143 ): 

144 super().__init__(bbox) 

145 # Set some useful parameters for derivations 

146 sin = np.sin(theta) 

147 cos = np.cos(theta) 

148 

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

150 # and xMinor as the y-axis 

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

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

153 # The scaled major and minor axes 

154 self._xa = self._xMajor / major 

155 self._yb = self._xMinor / minor 

156 

157 # Store parameters useful for gradient calculation 

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

159 self._theta = theta 

160 self._major = major 

161 self._minor = minor 

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

163 self._bbox = bbox 

164 self._rMin = r_min 

165 # Store the scaled radius**2 and radius 

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

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

168 

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

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

171 

172 Parameters 

173 ---------- 

174 input_grad: 

175 Gradient of the likelihood wrt the component model 

176 use_r2: 

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

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

179 

180 Returns 

181 ------- 

182 result: 

183 The gradient of the likelihood wrt x0. 

184 """ 

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

186 if use_r2: 

187 grad *= 2 

188 else: 

189 r = self.r_grid 

190 grad *= 1 / r 

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

192 

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

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

195 

196 Parameters 

197 ---------- 

198 input_grad: 

199 Gradient of the likelihood wrt the component model 

200 use_r2: 

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

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

203 

204 Returns 

205 ------- 

206 result: 

207 The gradient of the likelihood wrt y0. 

208 """ 

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

210 if use_r2: 

211 grad *= 2 

212 else: 

213 r = self.r_grid 

214 grad *= 1 / r 

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

216 

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

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

219 the semi-major axis 

220 

221 Parameters 

222 ---------- 

223 input_grad: 

224 Gradient of the likelihood wrt the component model 

225 use_r2: 

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

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

228 

229 Returns 

230 ------- 

231 result: 

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

233 """ 

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

235 if use_r2: 

236 grad *= 2 

237 else: 

238 r = self.r_grid 

239 grad *= 1 / r 

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

241 

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

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

244 the semi-minor axis 

245 

246 Parameters 

247 ---------- 

248 input_grad: 

249 Gradient of the likelihood wrt the component model 

250 use_r2: 

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

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

253 

254 Returns 

255 ------- 

256 result: 

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

258 """ 

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

260 if use_r2: 

261 grad *= 2 

262 else: 

263 r = self.r_grid 

264 grad *= 1 / r 

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

266 

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

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

269 the rotation angle 

270 

271 Parameters 

272 ---------- 

273 input_grad: 

274 Gradient of the likelihood wrt the component model 

275 use_r2: 

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

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

278 

279 Returns 

280 ------- 

281 result: 

282 The gradient of the likelihood wrt the rotation angle. 

283 """ 

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

285 if use_r2: 

286 grad *= 2 

287 else: 

288 r = self.r_grid 

289 grad *= 1 / r 

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

291 

292 @property 

293 def x0(self) -> float: 

294 """The x-center""" 

295 return self._x0 

296 

297 @property 

298 def y0(self) -> float: 

299 """The y-center""" 

300 return self._y0 

301 

302 @property 

303 def major(self) -> float: 

304 """The semi-major axis""" 

305 return self._major 

306 

307 @property 

308 def minor(self) -> float: 

309 """The semi-minor axis""" 

310 return self._minor 

311 

312 @property 

313 def theta(self) -> float: 

314 """The rotation angle""" 

315 return self._theta 

316 

317 @property 

318 def bbox(self) -> Box: 

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

320 return self._bbox 

321 

322 @property 

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

324 """The radial coordinates of each pixel""" 

325 if self._radius is None: 

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

327 return self._radius 

328 

329 @property 

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

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

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

333 

334 

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

336 """Model of a 2D elliptical gaussian 

337 

338 Parameters 

339 ---------- 

340 params: 

341 The parameters of the function. 

342 In this case there are none outside of the ellipticity 

343 ellipse: 

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

345 

346 Returns 

347 ------- 

348 result: 

349 The 2D guassian for the given ellipse parameters 

350 """ 

351 return np.exp(-ellipse.r2_grid) 

352 

353 

354def grad_gaussian2( 

355 input_grad: np.ndarray, 

356 params: np.ndarray, 

357 morph: np.ndarray, 

358 spectrum: np.ndarray, 

359 ellipse: EllipseFrame, 

360) -> np.ndarray: 

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

362 morphology parameters 

363 

364 Parameters 

365 ---------- 

366 input_grad: 

367 Gradient of the likelihood wrt the component model 

368 params: 

369 The parameters of the morphology. 

370 morph: 

371 The model of the morphology. 

372 spectrum: 

373 The model of the spectrum. 

374 ellipse: 

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

376 """ 

377 # Calculate the gradient of the likelihod 

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

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

380 d_y0 = ellipse.grad_y0(_grad, True) 

381 d_x0 = ellipse.grad_x0(_grad, True) 

382 d_sigma_y = ellipse.grad_major(_grad, True) 

383 d_sigma_x = ellipse.grad_minor(_grad, True) 

384 d_theta = ellipse.grad_theta(_grad, True) 

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

386 

387 

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

389 """Model of a circularly symmetric Gaussian 

390 

391 Parameters 

392 ---------- 

393 center: 

394 The center of the Gaussian. 

395 frame: 

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

397 sigma: 

398 The standard deviation. 

399 

400 Returns 

401 ------- 

402 result: 

403 The image of the circular Gaussian. 

404 """ 

405 y0, x0 = center[:2] 

406 two_sigma = 2 * sigma 

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

408 return np.exp(-r2) 

409 

410 

411def grad_circular_gaussian( 

412 input_grad: np.ndarray, 

413 params: np.ndarray, 

414 morph: np.ndarray, 

415 spectrum: np.ndarray, 

416 frame: CartesianFrame, 

417 sigma: float, 

418) -> np.ndarray: 

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

420 morphology parameters 

421 

422 Parameters 

423 ---------- 

424 input_grad: 

425 Gradient of the likelihood wrt the component model 

426 params: 

427 The parameters of the morphology. 

428 morph: 

429 The model of the morphology. 

430 spectrum: 

431 The model of the spectrum. 

432 frame: 

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

434 sigma: 

435 The standard deviation. 

436 """ 

437 # Calculate the gradient of the likelihod 

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

439 

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

441 

442 y0, x0 = params[:2] 

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

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

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

446 

447 

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

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

450 

451 This differs from `circularGaussian` because the gaussian function 

452 is integrated over each pixel to replicate the pixelated image 

453 version of a Gaussian function. 

454 

455 Parameters 

456 ---------- 

457 params: 

458 The center of the Gaussian. 

459 frame: 

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

461 

462 Returns 

463 ------- 

464 result: 

465 The image of the circular Gaussian. 

466 """ 

467 # Unpack the parameters and define constants 

468 y0, x0, sigma = params 

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

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

471 # Integrate from half a pixel left and right 

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

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

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

475 return z 

476 

477 

478def grad_integrated_gaussian( 

479 input_grad: np.ndarray, 

480 params: np.ndarray, 

481 morph: np.ndarray, 

482 spectrum: np.ndarray, 

483 frame: CartesianFrame, 

484) -> np.ndarray: 

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

486 morphology parameters 

487 

488 Parameters 

489 ---------- 

490 input_grad: 

491 Gradient of the likelihood wrt the component model parameters. 

492 params: 

493 The parameters of the morphology. 

494 morph: 

495 The model of the morphology. 

496 spectrum: 

497 The model of the spectrum. 

498 frame: 

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

500 

501 Returns 

502 ------- 

503 result: 

504 The gradient of the component morphology. 

505 """ 

506 # Calculate the gradient of the likelihood 

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

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

509 

510 # Extract the parameters 

511 y0, x0, sigma = params 

512 # define useful constants 

513 x = frame.x_grid - x0 

514 y = frame.y_grid - y0 

515 c = 0.5 / sigma**2 

516 sqrt_c = np.sqrt(c) 

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

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

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

520 r1 = r - 0.5 

521 r2 = r + 0.5 

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

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

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

525 # Calculate the gradients of the parameters 

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

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

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

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

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

531 

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

533 

534 

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

536 """A bounded proximal operator 

537 

538 This function updates `params` in place. 

539 

540 Parameters 

541 ---------- 

542 params: 

543 The array of parameters to constrain. 

544 proxmin: 

545 The array of minimum values for each parameter. 

546 proxmax: 

547 The array of maximum values for each parameter. 

548 

549 Returns 

550 ------- 

551 result: 

552 The updated parameters. 

553 """ 

554 cuts = params < proxmin 

555 params[cuts] = proxmin[cuts] 

556 cuts = params > proxmax 

557 params[cuts] = proxmax[cuts] 

558 return params 

559 

560 

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

562 """Generate a Sersic Model. 

563 

564 Parameters 

565 ---------- 

566 params: 

567 The parameters of the function. 

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

569 ellipse: 

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

571 

572 Returns 

573 ------- 

574 result: 

575 The model for the given ellipse parameters 

576 """ 

577 (n,) = params 

578 

579 r = ellipse.r_grid 

580 

581 if n == 1: 

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

583 else: 

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

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

586 return result 

587 

588 

589def grad_sersic( 

590 input_grad: np.ndarray, 

591 params: np.ndarray, 

592 morph: np.ndarray, 

593 spectrum: np.ndarray, 

594 ellipse: EllipseFrame, 

595) -> np.ndarray: 

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

597 

598 Parameters 

599 ---------- 

600 input_grad: 

601 Gradient of the likelihood wrt the component model 

602 params: 

603 The parameters of the morphology. 

604 morph: 

605 The model of the morphology. 

606 spectrum: 

607 The model of the spectrum. 

608 ellipse: 

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

610 """ 

611 n = params[5] 

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

613 if n == 1: 

614 # Use a simplified model for faster calculation 

615 d_exp = -SERSIC_B1 * morph 

616 else: 

617 r = ellipse.r_grid 

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

619 

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

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

622 _grad = _grad * d_exp 

623 d_y0 = ellipse.grad_y0(_grad, False) 

624 d_x0 = ellipse.grad_x0(_grad, False) 

625 d_sigma_y = ellipse.grad_major(_grad, False) 

626 d_sigma_x = ellipse.grad_minor(_grad, False) 

627 d_theta = ellipse.grad_theta(_grad, False) 

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

629 

630 

631class ParametricComponent(Component): 

632 """A parametric model of an astrophysical source 

633 

634 Parameters 

635 ---------- 

636 bands: 

637 The bands used in the model. 

638 bbox: 

639 The bounding box that holds the model. 

640 spectrum: 

641 The spectrum of the component. 

642 morph_params: 

643 The parameters of the morphology. 

644 morph_func: 

645 The function to generate the 2D morphology image 

646 based on `morphParams`. 

647 morph_grad: 

648 The function to calculate the gradient of the 

649 likelihood wrt the morphological parameters. 

650 morph_prox: 

651 The proximal operator for the morphology parameters. 

652 morph_step: 

653 The function that calculates the gradient of the 

654 morphological model. 

655 prox_spectrum: 

656 Proximal operator for the spectrum. 

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

658 operator `self.prox_spectrum` is used. 

659 floor: 

660 The minimum value of the spectrum, used to prevent 

661 divergences in the gradients. 

662 """ 

663 

664 def __init__( 

665 self, 

666 bands: tuple, 

667 bbox: Box, 

668 spectrum: Parameter | np.ndarray, 

669 morph_params: Parameter | np.ndarray, 

670 morph_func: Callable, 

671 morph_grad: Callable, 

672 morph_prox: Callable, 

673 morph_step: Callable | np.ndarray, 

674 prox_spectrum: Callable | None = None, 

675 floor: float = 1e-20, 

676 ): 

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

678 

679 self._spectrum = parameter(spectrum) 

680 self._params = parameter(morph_params) 

681 self._func = morph_func 

682 self._morph_grad = morph_grad 

683 self._morph_prox = morph_prox 

684 self._morph_step = morph_step 

685 self._bbox = bbox 

686 if prox_spectrum is None: 

687 self._prox_spectrum: Callable = self.prox_spectrum 

688 else: 

689 self._prox_spectrum = prox_spectrum 

690 self.floor = floor 

691 

692 @property 

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

694 """The center of the component""" 

695 return self.y0, self.x0 

696 

697 @property 

698 def y0(self) -> float: 

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

700 return self._params.x[0] 

701 

702 @property 

703 def x0(self) -> float: 

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

705 return self._params.x[1] 

706 

707 @property 

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

709 """The array of spectrum values""" 

710 return self._spectrum.x 

711 

712 @property 

713 def frame(self) -> CartesianFrame: 

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

715 return CartesianFrame(self._bbox) 

716 

717 @property 

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

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

720 return self._params.x 

721 

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

723 """The 2D image of the morphology 

724 

725 This callable generates an image of the morphology 

726 in the given frame. 

727 

728 Parameters 

729 ---------- 

730 frame: 

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

732 placed in. 

733 

734 Returns 

735 ------- 

736 result: 

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

738 """ 

739 if frame is None: 

740 frame = self.frame 

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

742 

743 @property 

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

745 """The morphological model""" 

746 return self._get_morph(frame) 

747 

748 @property 

749 def prox_morph(self) -> Callable: 

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

751 return self._morph_prox 

752 

753 @property 

754 def grad_morph(self) -> Callable: 

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

756 morphological model 

757 """ 

758 return self._morph_grad 

759 

760 @property 

761 def morph_step(self) -> Callable: 

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

763 morphological model 

764 """ 

765 return cast(Callable, self._morph_step) 

766 

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

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

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

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

771 

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

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

774 

775 Parameters 

776 ---------- 

777 spectrum: 

778 The spectrum of the model. 

779 """ 

780 # prevent divergent spectrum 

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

782 return spectrum 

783 

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

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

786 

787 Parameters 

788 ---------- 

789 input_grad: 

790 Gradient of the likelihood wrt the component model 

791 spectrum: 

792 The model of the spectrum. 

793 morph: 

794 The model of the morphology. 

795 

796 Returns 

797 ------- 

798 result: 

799 The gradient of the likelihood wrt. the spectrum. 

800 """ 

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

802 

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

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

805 

806 Parameters 

807 ---------- 

808 it: 

809 The current iteration of the optimizer. 

810 input_grad: 

811 Gradient of the likelihood wrt the component model 

812 """ 

813 spectrum = self.spectrum.copy() 

814 morph = self.morph 

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

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

817 

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

819 """Resize the box that contains the model 

820 

821 Not yet implemented, so for now the model box 

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

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

824 """ 

825 return False 

826 

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

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

829 # Update the spectrum and morph in place 

830 parameterization(self) 

831 # update the parameters 

832 self._spectrum.grad = self.grad_spectrum 

833 self._spectrum.prox = self.prox_spectrum 

834 self._params.grad = self.grad_morph 

835 self._params.prox = self.prox_morph 

836 

837 

838class EllipticalParametricComponent(ParametricComponent): 

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

840 

841 Parameters 

842 ---------- 

843 bands: 

844 The bands used in the model. 

845 bbox: 

846 The bounding box that holds this component model. 

847 spectrum: 

848 The spectrum of the component. 

849 morph_params: 

850 The parameters passed to `morph_func` to 

851 generate the morphology in image space. 

852 morph_func: 

853 The function to generate the morphology 

854 based on `morphParams`. 

855 morph_grad: 

856 The function to calculate the gradient of the 

857 likelihood wrt the morphological parameters. 

858 morph_prox: 

859 The proximal operator for the morphology parameters. 

860 prox_spectrum: 

861 Proximal operator for the spectrum. 

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

863 operator `self.prox_spectrum` is used. 

864 floor: 

865 The minimum value of the spectrum, used to prevent 

866 divergences in the gradients. 

867 """ 

868 

869 def __init__( 

870 self, 

871 bands: tuple, 

872 bbox: Box, 

873 spectrum: Parameter | np.ndarray, 

874 morph_params: Parameter | np.ndarray, 

875 morph_func: Callable, 

876 morph_grad: Callable, 

877 morph_prox: Callable, 

878 morph_step: Callable | np.ndarray, 

879 prox_spectrum: Callable | None = None, 

880 floor: float = 1e-20, 

881 ): 

882 super().__init__( 

883 bands=bands, 

884 bbox=bbox, 

885 spectrum=spectrum, 

886 morph_params=morph_params, 

887 morph_func=morph_func, 

888 morph_grad=morph_grad, 

889 morph_prox=morph_prox, 

890 morph_step=morph_step, 

891 prox_spectrum=prox_spectrum, 

892 floor=floor, 

893 ) 

894 

895 @property 

896 def semi_major(self) -> float: 

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

898 return self._params.x[2] 

899 

900 @property 

901 def semi_minor(self) -> float: 

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

903 return self._params.x[3] 

904 

905 @property 

906 def theta(self) -> float: 

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

908 x-axis. 

909 """ 

910 return self._params.x[4] 

911 

912 @property 

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

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

915 return self._params.x[:5] 

916 

917 @property 

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

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

920 return self._params.x[5:] 

921 

922 @property 

923 def frame(self) -> EllipseFrame: 

924 """The `EllipseFrame` that parameterizes the model""" 

925 return EllipseFrame(*self.ellipse_params, self._bbox) # type: ignore 

926 

927 @property 

928 def morph_prox(self) -> Callable: 

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

930 return self._morph_prox 

931 

932 @property 

933 def morph_grad(self) -> Callable: 

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

935 morphological model 

936 """ 

937 return self._morph_grad 

938 

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

940 """Update the component 

941 

942 Parameters 

943 ---------- 

944 it: 

945 The current iteration of the optimizer. 

946 input_grad: 

947 Gradient of the likelihood wrt the component model 

948 """ 

949 ellipse = self.frame 

950 spectrum = self.spectrum.copy() 

951 morph = self._func(self.radial_params, ellipse) 

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

953 self._params.update(it, input_grad, morph, spectrum, ellipse)