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

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 copy import deepcopy 

41from typing import TYPE_CHECKING, Any, Callable, Sequence, cast 

42 

43import numpy as np 

44from scipy.special import erf 

45from scipy.stats import gamma 

46 

47from ..bbox import Box 

48from ..component import Component 

49from ..image import Image 

50from ..parameters import Parameter, parameter 

51from ..utils import convert_indices 

52 

53if TYPE_CHECKING: 

54 from ..io import ScarletComponentBaseData 

55 

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

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

58MIN_RADIUS = 1e-20 

59 

60# Useful constants 

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

62 

63# Stored sersic constants 

64SERSIC_B1 = gamma.ppf(0.5, 2) 

65 

66 

67class CartesianFrame: 

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

69 

70 Parameters 

71 ---------- 

72 bbox: 

73 The bounding box that contains this frame. 

74 """ 

75 

76 def __init__(self, bbox: Box): 

77 # Store the new bounding box 

78 self._bbox = bbox 

79 # Get the range of x and y 

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

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

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

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

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

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

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

87 self._r = None 

88 self._r2 = None 

89 

90 @property 

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

92 """Shape of the frame.""" 

93 return self._bbox.shape 

94 

95 @property 

96 def bbox(self) -> Box: 

97 """Bounding box containing the frame""" 

98 return self._bbox 

99 

100 @property 

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

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

103 return self._x 

104 

105 @property 

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

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

108 return self._y 

109 

110 

111class EllipseFrame(CartesianFrame): 

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

113 

114 This frame is used to calculate the coordinates of the 

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

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

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

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

119 

120 Parameters 

121 ---------- 

122 y0: float 

123 The y-center of the ellipse. 

124 x0: float 

125 The x-center of the ellipse. 

126 major: float 

127 The length of the semi-major axis. 

128 minor: float 

129 The length of the semi-minor axis. 

130 theta: float 

131 The counter-clockwise rotation angle 

132 from the semi-major axis. 

133 bbox: Box 

134 The bounding box that contains the entire frame. 

135 r_min: float 

136 The minimum value of the radius. 

137 This is used to prevent divergences that occur 

138 when calculating the gradient at radius == 0. 

139 """ 

140 

141 def __init__( 

142 self, 

143 y0: float, 

144 x0: float, 

145 major: float, 

146 minor: float, 

147 theta: float, 

148 bbox: Box, 

149 r_min: float = 1e-20, 

150 ): 

151 super().__init__(bbox) 

152 # Set some useful parameters for derivations 

153 sin = np.sin(theta) 

154 cos = np.cos(theta) 

155 

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

157 # and xMinor as the y-axis 

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

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

160 # The scaled major and minor axes 

161 self._xa = self._xMajor / major 

162 self._yb = self._xMinor / minor 

163 

164 # Store parameters useful for gradient calculation 

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

166 self._theta = theta 

167 self._major = major 

168 self._minor = minor 

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

170 self._bbox = bbox 

171 self._rMin = r_min 

172 # Store the scaled radius**2 and radius 

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

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

175 

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

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

178 

179 Parameters 

180 ---------- 

181 input_grad: 

182 Gradient of the likelihood wrt the component model 

183 use_r2: 

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

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

186 

187 Returns 

188 ------- 

189 result: 

190 The gradient of the likelihood wrt x0. 

191 """ 

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

193 if use_r2: 

194 grad *= 2 

195 else: 

196 r = self.r_grid 

197 grad *= 1 / r 

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

199 

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

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

202 

203 Parameters 

204 ---------- 

205 input_grad: 

206 Gradient of the likelihood wrt the component model 

207 use_r2: 

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

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

210 

211 Returns 

212 ------- 

213 result: 

214 The gradient of the likelihood wrt y0. 

215 """ 

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

217 if use_r2: 

218 grad *= 2 

219 else: 

220 r = self.r_grid 

221 grad *= 1 / r 

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

223 

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

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

226 the semi-major axis 

227 

228 Parameters 

229 ---------- 

230 input_grad: 

231 Gradient of the likelihood wrt the component model 

232 use_r2: 

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

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

235 

236 Returns 

237 ------- 

238 result: 

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

240 """ 

241 # ``grad`` here is half of d(r**2)/d(major), matching the 

242 # convention used by ``grad_x0`` / ``grad_y0`` / ``grad_theta`` 

243 # so the post-multiplications below produce d(r**2) and d(r) 

244 # respectively. 

245 grad = -1 / self._major * self._xa**2 

246 if use_r2: 

247 grad *= 2 

248 else: 

249 r = self.r_grid 

250 grad *= 1 / r 

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

252 

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

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

255 the semi-minor axis 

256 

257 Parameters 

258 ---------- 

259 input_grad: 

260 Gradient of the likelihood wrt the component model 

261 use_r2: 

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

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

264 

265 Returns 

266 ------- 

267 result: 

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

269 """ 

270 # See ``grad_major`` for the half-r**2 convention. 

271 grad = -1 / self._minor * self._yb**2 

272 if use_r2: 

273 grad *= 2 

274 else: 

275 r = self.r_grid 

276 grad *= 1 / r 

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

278 

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

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

281 the rotation angle 

282 

283 Parameters 

284 ---------- 

285 input_grad: 

286 Gradient of the likelihood wrt the component model 

287 use_r2: 

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

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

290 

291 Returns 

292 ------- 

293 result: 

294 The gradient of the likelihood wrt the rotation angle. 

295 """ 

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

297 if use_r2: 

298 grad *= 2 

299 else: 

300 r = self.r_grid 

301 grad *= 1 / r 

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

303 

304 @property 

305 def x0(self) -> float: 

306 """The x-center""" 

307 return self._x0 

308 

309 @property 

310 def y0(self) -> float: 

311 """The y-center""" 

312 return self._y0 

313 

314 @property 

315 def major(self) -> float: 

316 """The semi-major axis""" 

317 return self._major 

318 

319 @property 

320 def minor(self) -> float: 

321 """The semi-minor axis""" 

322 return self._minor 

323 

324 @property 

325 def theta(self) -> float: 

326 """The rotation angle""" 

327 return self._theta 

328 

329 @property 

330 def bbox(self) -> Box: 

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

332 return self._bbox 

333 

334 @property 

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

336 """The radial coordinates of each pixel""" 

337 if self._radius is None: 

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

339 return self._radius 

340 

341 @property 

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

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

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

345 

346 

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

348 """Model of a 2D elliptical gaussian 

349 

350 Parameters 

351 ---------- 

352 params: 

353 The parameters of the function. 

354 In this case there are none outside of the ellipticity 

355 ellipse: 

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

357 

358 Returns 

359 ------- 

360 result: 

361 The 2D guassian for the given ellipse parameters 

362 """ 

363 return np.exp(-ellipse.r2_grid) 

364 

365 

366def grad_gaussian2( 

367 input_grad: np.ndarray, 

368 params: np.ndarray, 

369 morph: np.ndarray, 

370 spectrum: np.ndarray, 

371 ellipse: EllipseFrame, 

372) -> np.ndarray: 

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

374 morphology parameters 

375 

376 Parameters 

377 ---------- 

378 input_grad: 

379 Gradient of the likelihood wrt the component model 

380 params: 

381 The parameters of the morphology. 

382 morph: 

383 The model of the morphology. 

384 spectrum: 

385 The model of the spectrum. 

386 ellipse: 

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

388 """ 

389 # Calculate the gradient of the likelihod 

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

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

392 d_y0 = ellipse.grad_y0(_grad, True) 

393 d_x0 = ellipse.grad_x0(_grad, True) 

394 d_sigma_y = ellipse.grad_major(_grad, True) 

395 d_sigma_x = ellipse.grad_minor(_grad, True) 

396 d_theta = ellipse.grad_theta(_grad, True) 

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

398 

399 

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

401 """Model of a circularly symmetric Gaussian 

402 

403 Parameters 

404 ---------- 

405 center: 

406 The center of the Gaussian. 

407 frame: 

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

409 sigma: 

410 The standard deviation. 

411 

412 Returns 

413 ------- 

414 result: 

415 The image of the circular Gaussian. 

416 """ 

417 y0, x0 = center[:2] 

418 two_sigma = 2 * sigma 

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

420 return np.exp(-r2) 

421 

422 

423def grad_circular_gaussian( 

424 input_grad: np.ndarray, 

425 params: np.ndarray, 

426 morph: np.ndarray, 

427 spectrum: np.ndarray, 

428 frame: CartesianFrame, 

429 sigma: float, 

430) -> np.ndarray: 

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

432 morphology parameters 

433 

434 Parameters 

435 ---------- 

436 input_grad: 

437 Gradient of the likelihood wrt the component model 

438 params: 

439 The parameters of the morphology. 

440 morph: 

441 The model of the morphology. 

442 spectrum: 

443 The model of the spectrum. 

444 frame: 

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

446 sigma: 

447 The standard deviation. 

448 """ 

449 # Calculate the gradient of the likelihod 

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

451 

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

453 

454 y0, x0 = params[:2] 

455 # d morph / d y0 = morph * (y - y0) / (2 * sigma**2), and similarly for x0, 

456 # because r2 = ((x-x0)/(2*sigma))**2 + ((y-y0)/(2*sigma))**2. 

457 inv_two_sigma_sq = 1.0 / (2.0 * sigma**2) 

458 d_y0 = -inv_two_sigma_sq * np.sum((frame.y_grid - y0) * _grad) 

459 d_x0 = -inv_two_sigma_sq * np.sum((frame.x_grid - x0) * _grad) 

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

461 

462 

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

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

465 

466 This differs from `circularGaussian` because the gaussian function 

467 is integrated over each pixel to replicate the pixelated image 

468 version of a Gaussian function. 

469 

470 Parameters 

471 ---------- 

472 params: 

473 The center of the Gaussian. 

474 frame: 

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

476 

477 Returns 

478 ------- 

479 result: 

480 The image of the circular Gaussian. 

481 """ 

482 # Unpack the parameters and define constants 

483 y0, x0, sigma = params 

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

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

486 # Integrate from half a pixel left and right 

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

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

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

490 return z 

491 

492 

493def grad_integrated_gaussian( 

494 input_grad: np.ndarray, 

495 params: np.ndarray, 

496 morph: np.ndarray, 

497 spectrum: np.ndarray, 

498 frame: CartesianFrame, 

499) -> np.ndarray: 

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

501 morphology parameters 

502 

503 Parameters 

504 ---------- 

505 input_grad: 

506 Gradient of the likelihood wrt the component model parameters. 

507 params: 

508 The parameters of the morphology. 

509 morph: 

510 The model of the morphology. 

511 spectrum: 

512 The model of the spectrum. 

513 frame: 

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

515 

516 Returns 

517 ------- 

518 result: 

519 The gradient of the component morphology. 

520 """ 

521 # Calculate the gradient of the likelihood 

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

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

524 

525 # Extract the parameters 

526 y0, x0, sigma = params 

527 # define useful constants 

528 x = frame.x_grid - x0 

529 y = frame.y_grid - y0 

530 c = 0.5 / sigma**2 

531 sqrt_c = np.sqrt(c) 

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

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

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

535 r1 = r - 0.5 

536 r2 = r + 0.5 

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

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

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

540 # Calculate the gradients of the parameters 

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

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

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

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

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

546 

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

548 

549 

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

551 """A bounded proximal operator 

552 

553 This function updates `params` in place. 

554 

555 Parameters 

556 ---------- 

557 params: 

558 The array of parameters to constrain. 

559 proxmin: 

560 The array of minimum values for each parameter. 

561 proxmax: 

562 The array of maximum values for each parameter. 

563 

564 Returns 

565 ------- 

566 result: 

567 The updated parameters. 

568 """ 

569 cuts = params < proxmin 

570 params[cuts] = proxmin[cuts] 

571 cuts = params > proxmax 

572 params[cuts] = proxmax[cuts] 

573 return params 

574 

575 

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

577 """Generate a Sersic Model. 

578 

579 Parameters 

580 ---------- 

581 params: 

582 The parameters of the function. 

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

584 ellipse: 

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

586 

587 Returns 

588 ------- 

589 result: 

590 The model for the given ellipse parameters 

591 """ 

592 (n,) = params 

593 

594 r = ellipse.r_grid 

595 

596 if n == 1: 

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

598 else: 

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

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

601 return result 

602 

603 

604def grad_sersic( 

605 input_grad: np.ndarray, 

606 params: np.ndarray, 

607 morph: np.ndarray, 

608 spectrum: np.ndarray, 

609 ellipse: EllipseFrame, 

610) -> np.ndarray: 

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

612 

613 Parameters 

614 ---------- 

615 input_grad: 

616 Gradient of the likelihood wrt the component model 

617 params: 

618 The parameters of the morphology. 

619 morph: 

620 The model of the morphology. 

621 spectrum: 

622 The model of the spectrum. 

623 ellipse: 

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

625 """ 

626 n = params[5] 

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

628 if n == 1: 

629 # Use a simplified model for faster calculation 

630 d_exp = -SERSIC_B1 * morph 

631 else: 

632 r = ellipse.r_grid 

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

634 

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

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

637 _grad = _grad * d_exp 

638 d_y0 = ellipse.grad_y0(_grad, False) 

639 d_x0 = ellipse.grad_x0(_grad, False) 

640 d_sigma_y = ellipse.grad_major(_grad, False) 

641 d_sigma_x = ellipse.grad_minor(_grad, False) 

642 d_theta = ellipse.grad_theta(_grad, False) 

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

644 

645 

646class ParametricComponent(Component): 

647 """A parametric model of an astrophysical source 

648 

649 Parameters 

650 ---------- 

651 bands: 

652 The bands used in the model. 

653 bbox: 

654 The bounding box that holds the model. 

655 spectrum: 

656 The spectrum of the component. 

657 morph_params: 

658 The parameters of the morphology. 

659 morph_func: 

660 The function to generate the 2D morphology image 

661 based on `morphParams`. 

662 morph_grad: 

663 The function to calculate the gradient of the 

664 likelihood wrt the morphological parameters. 

665 morph_prox: 

666 The proximal operator for the morphology parameters. 

667 morph_step: 

668 The function that calculates the gradient of the 

669 morphological model. 

670 prox_spectrum: 

671 Proximal operator for the spectrum. 

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

673 operator `self.prox_spectrum` is used. 

674 floor: 

675 The minimum value of the spectrum, used to prevent 

676 divergences in the gradients. 

677 """ 

678 

679 def __init__( 

680 self, 

681 bands: tuple, 

682 bbox: Box, 

683 spectrum: Parameter | np.ndarray, 

684 morph_params: Parameter | np.ndarray, 

685 morph_func: Callable, 

686 morph_grad: Callable, 

687 morph_prox: Callable, 

688 morph_step: Callable | np.ndarray, 

689 prox_spectrum: Callable | None = None, 

690 floor: float = 1e-20, 

691 ): 

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

693 

694 self._spectrum = parameter(spectrum) 

695 self._params = parameter(morph_params) 

696 self._func = morph_func 

697 self._morph_grad = morph_grad 

698 self._morph_prox = morph_prox 

699 self._morph_step = morph_step 

700 self._bbox = bbox 

701 if prox_spectrum is None: 

702 self._prox_spectrum: Callable = self.prox_spectrum 

703 else: 

704 self._prox_spectrum = prox_spectrum 

705 self.floor = floor 

706 

707 @property 

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

709 """The center of the component""" 

710 return self.y0, self.x0 

711 

712 @property 

713 def y0(self) -> float: 

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

715 return self._params.x[0] 

716 

717 @property 

718 def x0(self) -> float: 

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

720 return self._params.x[1] 

721 

722 @property 

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

724 """The array of spectrum values""" 

725 return self._spectrum.x 

726 

727 @property 

728 def frame(self) -> CartesianFrame: 

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

730 return CartesianFrame(self._bbox) 

731 

732 @property 

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

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

735 return self._params.x 

736 

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

738 """The 2D image of the morphology 

739 

740 This callable generates an image of the morphology 

741 in the given frame. 

742 

743 Parameters 

744 ---------- 

745 frame: 

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

747 placed in. 

748 

749 Returns 

750 ------- 

751 result: 

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

753 """ 

754 if frame is None: 

755 frame = self.frame 

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

757 

758 @property 

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

760 """The morphological model""" 

761 return self._get_morph(frame) 

762 

763 @property 

764 def prox_morph(self) -> Callable: 

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

766 return self._morph_prox 

767 

768 @property 

769 def grad_morph(self) -> Callable: 

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

771 morphological model 

772 """ 

773 return self._morph_grad 

774 

775 @property 

776 def morph_step(self) -> Callable: 

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

778 morphological model 

779 """ 

780 return cast(Callable, self._morph_step) 

781 

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

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

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

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

786 

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

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

789 

790 Parameters 

791 ---------- 

792 spectrum: 

793 The spectrum of the model. 

794 """ 

795 # prevent divergent spectrum 

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

797 return spectrum 

798 

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

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

801 

802 Parameters 

803 ---------- 

804 input_grad: 

805 Gradient of the likelihood wrt the component model 

806 spectrum: 

807 The model of the spectrum. 

808 morph: 

809 The model of the morphology. 

810 

811 Returns 

812 ------- 

813 result: 

814 The gradient of the likelihood wrt. the spectrum. 

815 """ 

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

817 

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

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

820 

821 Parameters 

822 ---------- 

823 it: 

824 The current iteration of the optimizer. 

825 input_grad: 

826 Gradient of the likelihood wrt the component model 

827 """ 

828 spectrum = self.spectrum.copy() 

829 morph = self.morph 

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

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

832 

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

834 """Resize the box that contains the model 

835 

836 Not yet implemented, so for now the model box 

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

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

839 """ 

840 return False 

841 

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

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

844 # Update the spectrum and morph in place 

845 parameterization(self) 

846 # update the parameters 

847 self._spectrum.grad = self.grad_spectrum 

848 self._spectrum.prox = self.prox_spectrum 

849 self._params.grad = self.grad_morph 

850 self._params.prox = self.prox_morph 

851 

852 def to_data(self) -> ScarletComponentBaseData: 

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

854 

855 def __getitem__(self, indices: Any) -> ParametricComponent: 

856 """Get a sub-component corresponding to the given indices. 

857 

858 Parameters 

859 ---------- 

860 indices: Any 

861 The indices to use to slice the component model. 

862 

863 Returns 

864 ------- 

865 component: ParametricComponent 

866 A new component that is a sub-component of this one. 

867 

868 Raises 

869 ------ 

870 IndexError : 

871 If the index includes a ``Box`` or spatial indices. 

872 """ 

873 # Update the bands 

874 if indices in self._bands: 

875 # Single band case 

876 bands = (indices,) 

877 else: 

878 # Multiple bands case 

879 bands = tuple(indices) 

880 

881 # Convert the band indices into numerical indices 

882 band_indices = convert_indices(self.bands, indices) 

883 

884 # Slice the spectrum 

885 spectrum = self._spectrum.x[band_indices] 

886 

887 return ParametricComponent( 

888 bands=bands, 

889 bbox=self.bbox, 

890 spectrum=spectrum, 

891 morph_params=self.radial_params, 

892 morph_func=self._func, 

893 morph_grad=self._morph_grad, 

894 morph_prox=self._morph_prox, 

895 morph_step=self._morph_step, 

896 prox_spectrum=self._prox_spectrum, 

897 floor=self.floor, 

898 ) 

899 

900 def __deepcopy__(self, memo: dict[int, Any]) -> ParametricComponent: 

901 """Create a deep copy of this component 

902 

903 Parameters 

904 ---------- 

905 memo: 

906 The memoization dictionary used by `copy.deepcopy`. 

907 Returns 

908 ------- 

909 component : ParametricComponent 

910 A new component that is a deep copy of this one. 

911 """ 

912 if id(self) in memo: 

913 return memo[id(self)] 

914 

915 component = ParametricComponent.__new__(ParametricComponent) 

916 memo[id(self)] = component 

917 

918 component.__init__( # type: ignore[misc] 

919 bands=deepcopy(self.bands), 

920 bbox=deepcopy(self.bbox), 

921 spectrum=deepcopy(self.spectrum), 

922 morph_params=deepcopy(self.radial_params), 

923 morph_func=self._func, 

924 morph_grad=self._morph_grad, 

925 morph_prox=self._morph_prox, 

926 morph_step=self._morph_step, 

927 prox_spectrum=self._prox_spectrum,