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 def __init__(self, bbox: Box): 

71 """ 

72 Parameters 

73 ---------- 

74 bbox: Box 

75 The bounding box that contains this frame. 

76 """ 

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 = -2 / self._major * self._xa**2 

242 if use_r2: 

243 grad *= 2 

244 else: 

245 r = self.r_grid 

246 grad *= 1 / r 

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

248 

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

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

251 the semi-minor axis 

252 

253 Parameters 

254 ---------- 

255 input_grad: 

256 Gradient of the likelihood wrt the component model 

257 use_r2: 

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

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

260 

261 Returns 

262 ------- 

263 result: 

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

265 """ 

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

267 if use_r2: 

268 grad *= 2 

269 else: 

270 r = self.r_grid 

271 grad *= 1 / r 

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

273 

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

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

276 the rotation angle 

277 

278 Parameters 

279 ---------- 

280 input_grad: 

281 Gradient of the likelihood wrt the component model 

282 use_r2: 

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

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

285 

286 Returns 

287 ------- 

288 result: 

289 The gradient of the likelihood wrt the rotation angle. 

290 """ 

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

292 if use_r2: 

293 grad *= 2 

294 else: 

295 r = self.r_grid 

296 grad *= 1 / r 

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

298 

299 @property 

300 def x0(self) -> float: 

301 """The x-center""" 

302 return self._x0 

303 

304 @property 

305 def y0(self) -> float: 

306 """The y-center""" 

307 return self._y0 

308 

309 @property 

310 def major(self) -> float: 

311 """The semi-major axis""" 

312 return self._major 

313 

314 @property 

315 def minor(self) -> float: 

316 """The semi-minor axis""" 

317 return self._minor 

318 

319 @property 

320 def theta(self) -> float: 

321 """The rotation angle""" 

322 return self._theta 

323 

324 @property 

325 def bbox(self) -> Box: 

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

327 return self._bbox 

328 

329 @property 

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

331 """The radial coordinates of each pixel""" 

332 if self._radius is None: 

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

334 return self._radius 

335 

336 @property 

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

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

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

340 

341 

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

343 """Model of a 2D elliptical gaussian 

344 

345 Parameters 

346 ---------- 

347 params: 

348 The parameters of the function. 

349 In this case there are none outside of the ellipticity 

350 ellipse: 

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

352 

353 Returns 

354 ------- 

355 result: 

356 The 2D guassian for the given ellipse parameters 

357 """ 

358 return np.exp(-ellipse.r2_grid) 

359 

360 

361def grad_gaussian2( 

362 input_grad: np.ndarray, 

363 params: np.ndarray, 

364 morph: np.ndarray, 

365 spectrum: np.ndarray, 

366 ellipse: EllipseFrame, 

367) -> np.ndarray: 

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

369 morphology parameters 

370 

371 Parameters 

372 ---------- 

373 input_grad: 

374 Gradient of the likelihood wrt the component model 

375 params: 

376 The parameters of the morphology. 

377 morph: 

378 The model of the morphology. 

379 spectrum: 

380 The model of the spectrum. 

381 ellipse: 

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

383 """ 

384 # Calculate the gradient of the likelihod 

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

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

387 d_y0 = ellipse.grad_y0(_grad, True) 

388 d_x0 = ellipse.grad_x0(_grad, True) 

389 d_sigma_y = ellipse.grad_major(_grad, True) 

390 d_sigma_x = ellipse.grad_minor(_grad, True) 

391 d_theta = ellipse.grad_theta(_grad, True) 

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

393 

394 

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

396 """Model of a circularly symmetric Gaussian 

397 

398 Parameters 

399 ---------- 

400 center: 

401 The center of the Gaussian. 

402 frame: 

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

404 sigma: 

405 The standard deviation. 

406 

407 Returns 

408 ------- 

409 result: 

410 The image of the circular Gaussian. 

411 """ 

412 y0, x0 = center[:2] 

413 two_sigma = 2 * sigma 

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

415 return np.exp(-r2) 

416 

417 

418def grad_circular_gaussian( 

419 input_grad: np.ndarray, 

420 params: np.ndarray, 

421 morph: np.ndarray, 

422 spectrum: np.ndarray, 

423 frame: CartesianFrame, 

424 sigma: float, 

425) -> np.ndarray: 

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

427 morphology parameters 

428 

429 Parameters 

430 ---------- 

431 input_grad: 

432 Gradient of the likelihood wrt the component model 

433 params: 

434 The parameters of the morphology. 

435 morph: 

436 The model of the morphology. 

437 spectrum: 

438 The model of the spectrum. 

439 frame: 

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

441 sigma: 

442 The standard deviation. 

443 """ 

444 # Calculate the gradient of the likelihod 

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

446 

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

448 

449 y0, x0 = params[:2] 

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

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

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

453 

454 

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

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

457 

458 This differs from `circularGaussian` because the gaussian function 

459 is integrated over each pixel to replicate the pixelated image 

460 version of a Gaussian function. 

461 

462 Parameters 

463 ---------- 

464 params: 

465 The center of the Gaussian. 

466 frame: 

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

468 

469 Returns 

470 ------- 

471 result: 

472 The image of the circular Gaussian. 

473 """ 

474 # Unpack the parameters and define constants 

475 y0, x0, sigma = params 

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

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

478 # Integrate from half a pixel left and right 

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

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

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

482 return z 

483 

484 

485def grad_integrated_gaussian( 

486 input_grad: np.ndarray, 

487 params: np.ndarray, 

488 morph: np.ndarray, 

489 spectrum: np.ndarray, 

490 frame: CartesianFrame, 

491) -> np.ndarray: 

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

493 morphology parameters 

494 

495 Parameters 

496 ---------- 

497 input_grad: 

498 Gradient of the likelihood wrt the component model parameters. 

499 params: 

500 The parameters of the morphology. 

501 morph: 

502 The model of the morphology. 

503 spectrum: 

504 The model of the spectrum. 

505 frame: 

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

507 

508 Returns 

509 ------- 

510 result: 

511 The gradient of the component morphology. 

512 """ 

513 # Calculate the gradient of the likelihood 

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

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

516 

517 # Extract the parameters 

518 y0, x0, sigma = params 

519 # define useful constants 

520 x = frame.x_grid - x0 

521 y = frame.y_grid - y0 

522 c = 0.5 / sigma**2 

523 sqrt_c = np.sqrt(c) 

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

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

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

527 r1 = r - 0.5 

528 r2 = r + 0.5 

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

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

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

532 # Calculate the gradients of the parameters 

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

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

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

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

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

538 

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

540 

541 

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

543 """A bounded proximal operator 

544 

545 This function updates `params` in place. 

546 

547 Parameters 

548 ---------- 

549 params: 

550 The array of parameters to constrain. 

551 proxmin: 

552 The array of minimum values for each parameter. 

553 proxmax: 

554 The array of maximum values for each parameter. 

555 

556 Returns 

557 ------- 

558 result: 

559 The updated parameters. 

560 """ 

561 cuts = params < proxmin 

562 params[cuts] = proxmin[cuts] 

563 cuts = params > proxmax 

564 params[cuts] = proxmax[cuts] 

565 return params 

566 

567 

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

569 """Generate a Sersic Model. 

570 

571 Parameters 

572 ---------- 

573 params: 

574 The parameters of the function. 

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

576 ellipse: 

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

578 

579 Returns 

580 ------- 

581 result: 

582 The model for the given ellipse parameters 

583 """ 

584 (n,) = params 

585 

586 r = ellipse.r_grid 

587 

588 if n == 1: 

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

590 else: 

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

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

593 return result 

594 

595 

596def grad_sersic( 

597 input_grad: np.ndarray, 

598 params: np.ndarray, 

599 morph: np.ndarray, 

600 spectrum: np.ndarray, 

601 ellipse: EllipseFrame, 

602) -> np.ndarray: 

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

604 

605 Parameters 

606 ---------- 

607 input_grad: 

608 Gradient of the likelihood wrt the component model 

609 params: 

610 The parameters of the morphology. 

611 morph: 

612 The model of the morphology. 

613 spectrum: 

614 The model of the spectrum. 

615 ellipse: 

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

617 """ 

618 n = params[5] 

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

620 if n == 1: 

621 # Use a simplified model for faster calculation 

622 d_exp = -SERSIC_B1 * morph 

623 else: 

624 r = ellipse.r_grid 

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

626 

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

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

629 _grad = _grad * d_exp 

630 d_y0 = ellipse.grad_y0(_grad, False) 

631 d_x0 = ellipse.grad_x0(_grad, False) 

632 d_sigma_y = ellipse.grad_major(_grad, False) 

633 d_sigma_x = ellipse.grad_minor(_grad, False) 

634 d_theta = ellipse.grad_theta(_grad, False) 

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

636 

637 

638class ParametricComponent(Component): 

639 """A parametric model of an astrophysical source 

640 

641 Parameters 

642 ---------- 

643 bands: 

644 The bands used in the model. 

645 bbox: 

646 The bounding box that holds the model. 

647 spectrum: 

648 The spectrum of the component. 

649 morph_params: 

650 The parameters of the morphology. 

651 morph_func: 

652 The function to generate the 2D morphology image 

653 based on `morphParams`. 

654 morph_grad: 

655 The function to calculate the gradient of the 

656 likelihood wrt the morphological parameters. 

657 morph_prox: 

658 The proximal operator for the morphology parameters. 

659 morph_step: 

660 The function that calculates the gradient of the 

661 morphological model. 

662 prox_spectrum: 

663 Proximal operator for the spectrum. 

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

665 operator `self.prox_spectrum` is used. 

666 floor: 

667 The minimum value of the spectrum, used to prevent 

668 divergences in the gradients. 

669 """ 

670 

671 def __init__( 

672 self, 

673 bands: tuple, 

674 bbox: Box, 

675 spectrum: Parameter | np.ndarray, 

676 morph_params: Parameter | np.ndarray, 

677 morph_func: Callable, 

678 morph_grad: Callable, 

679 morph_prox: Callable, 

680 morph_step: Callable | np.ndarray, 

681 prox_spectrum: Callable | None = None, 

682 floor: float = 1e-20, 

683 ): 

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

685 

686 self._spectrum = parameter(spectrum) 

687 self._params = parameter(morph_params) 

688 self._func = morph_func 

689 self._morph_grad = morph_grad 

690 self._morph_prox = morph_prox 

691 self._morph_step = morph_step 

692 self._bbox = bbox 

693 if prox_spectrum is None: 

694 self._prox_spectrum: Callable = self.prox_spectrum 

695 else: 

696 self._prox_spectrum = prox_spectrum 

697 self.floor = floor 

698 

699 @property 

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

701 """The center of the component""" 

702 return self.y0, self.x0 

703 

704 @property 

705 def y0(self) -> float: 

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

707 return self._params.x[0] 

708 

709 @property 

710 def x0(self) -> float: 

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

712 return self._params.x[1] 

713 

714 @property 

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

716 """The array of spectrum values""" 

717 return self._spectrum.x 

718 

719 @property 

720 def frame(self) -> CartesianFrame: 

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

722 return CartesianFrame(self._bbox) 

723 

724 @property 

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

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

727 return self._params.x 

728 

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

730 """The 2D image of the morphology 

731 

732 This callable generates an image of the morphology 

733 in the given frame. 

734 

735 Parameters 

736 ---------- 

737 frame: 

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

739 placed in. 

740 

741 Returns 

742 ------- 

743 result: 

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

745 """ 

746 if frame is None: 

747 frame = self.frame 

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

749 

750 @property 

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

752 """The morphological model""" 

753 return self._get_morph(frame) 

754 

755 @property 

756 def prox_morph(self) -> Callable: 

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

758 return self._morph_prox 

759 

760 @property 

761 def grad_morph(self) -> Callable: 

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

763 morphological model 

764 """ 

765 return self._morph_grad 

766 

767 @property 

768 def morph_step(self) -> Callable: 

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

770 morphological model 

771 """ 

772 return cast(Callable, self._morph_step) 

773 

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

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

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

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

778 

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

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

781 

782 Parameters 

783 ---------- 

784 spectrum: 

785 The spectrum of the model. 

786 """ 

787 # prevent divergent spectrum 

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

789 return spectrum 

790 

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

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

793 

794 Parameters 

795 ---------- 

796 input_grad: 

797 Gradient of the likelihood wrt the component model 

798 spectrum: 

799 The model of the spectrum. 

800 morph: 

801 The model of the morphology. 

802 

803 Returns 

804 ------- 

805 result: 

806 The gradient of the likelihood wrt. the spectrum. 

807 """ 

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

809 

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

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

812 

813 Parameters 

814 ---------- 

815 it: 

816 The current iteration of the optimizer. 

817 input_grad: 

818 Gradient of the likelihood wrt the component model 

819 """ 

820 spectrum = self.spectrum.copy() 

821 morph = self.morph 

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

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

824 

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

826 """Resize the box that contains the model 

827 

828 Not yet implemented, so for now the model box 

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

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

831 """ 

832 return False 

833 

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

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

836 # Update the spectrum and morph in place 

837 parameterization(self) 

838 # update the parameters 

839 self._spectrum.grad = self.grad_spectrum 

840 self._spectrum.prox = self.prox_spectrum 

841 self._params.grad = self.grad_morph 

842 self._params.prox = self.prox_morph 

843 

844 def to_data(self) -> ScarletComponentBaseData: 

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

846 

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

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

849 

850 Parameters 

851 ---------- 

852 indices: Any 

853 The indices to use to slice the component model. 

854 

855 Returns 

856 ------- 

857 component: ParametricComponent 

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

859 

860 Raises 

861 ------ 

862 IndexError : 

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

864 """ 

865 # Update the bands 

866 if indices in self._bands: 

867 # Single band case 

868 bands = (indices,) 

869 else: 

870 # Multiple bands case 

871 bands = tuple(indices) 

872 

873 # Convert the band indices into numerical indices 

874 band_indices = convert_indices(self.bands, indices) 

875 

876 # Slice the spectrum 

877 spectrum = self._spectrum.x[band_indices] 

878 

879 return ParametricComponent( 

880 bands=bands, 

881 bbox=self.bbox, 

882 spectrum=spectrum, 

883 morph_params=self.radial_params, 

884 morph_func=self._func, 

885 morph_grad=self._morph_grad, 

886 morph_prox=self._morph_prox, 

887 morph_step=self._morph_step, 

888 prox_spectrum=self._prox_spectrum, 

889 floor=self.floor, 

890 ) 

891 

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

893 """Create a deep copy of this component 

894 

895 Parameters 

896 ---------- 

897 memo: 

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

899 Returns 

900 ------- 

901 component : ParametricComponent 

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

903 """ 

904 if id(self) in memo: 

905 return memo[id(self)] 

906 

907 component = ParametricComponent.__new__(ParametricComponent) 

908 memo[id(self)] = component 

909 

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

911 bands=deepcopy(self.bands), 

912 bbox=deepcopy(self.bbox), 

913 spectrum=deepcopy(self.spectrum), 

914 morph_params=deepcopy(self.radial_params), 

915 morph_func=self._func, 

916 morph_grad=self._morph_grad, 

917 morph_prox=self._morph_prox, 

918 morph_step=self._morph_step, 

919 prox_spectrum=self._prox_spectrum, 

920 floor=self.floor, 

921 ) 

922 

923 return component 

924 

925 def __copy__(self) -> ParametricComponent: 

926 """Create a copy of this component