Coverage for tests / utils.py: 26%

1# This file is part of meas_extensions_scarlet. 

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 typing import Sequence 

23 

24import lsst.scarlet.lite as scl 

25import numpy as np 

26import lsst.meas.extensions.scarlet as mes 

27 

28 

29class DeblenderTestModel: 

30 center: tuple[float, float] 

31 spectrum: np.ndarray 

32 morph: np.ndarray 

33 bbox: scl.Box 

34 bands: Sequence[str] 

35 

36 def render(self, psf: np.ndarray) -> scl.Image: 

37 model = self.spectrum[:, None, None]*self.morph[None, :, :] 

38 if len(psf.shape) == 2: 

39 psf = np.repeat(psf[None, :, :], model.shape[0], axis=0) 

40 model = mes.utils.multiband_convolve(model, psf) 

41 return scl.Image(model, bands=self.bands, yx0=self.bbox.origin) 

42 

43 

44class SersicModel(DeblenderTestModel): 

45 def __init__( 

46 self, 

47 center: tuple[int, int], 

48 major: float, 

49 minor: float, 

50 radius: int, 

51 theta: float, 

52 n: float, 

53 spectrum: np.ndarray, 

54 bands: Sequence[str], 

55 ): 

56 self.center = center 

57 self.spectrum = spectrum.astype(np.float32) 

58 self.bands = bands 

59 bbox = scl.Box((2*radius+1, 2*radius+1), (center[0]-radius, center[1]-radius)) 

60 self.bbox = bbox 

61 frame = scl.models.parametric.EllipseFrame( 

62 center[0], 

63 center[1], 

64 major, 

65 minor, 

66 theta, 

67 bbox, 

68 ) 

69 morph = scl.models.parametric.sersic((n,), frame).astype(np.float32) 

70 self.morph = morph/np.max(morph) 

71 

72 

73class PsfModel(DeblenderTestModel): 

74 def __init__( 

75 self, 

76 center, 

77 spectrum: np.ndarray, 

78 bands: Sequence[str], 

79 radius: int = 7, 

80 ): 

81 self.center = center 

82 self.spectrum = spectrum.astype(np.float32) 

83 self.bands = bands 

84 bbox = scl.Box((2*radius+1, 2*radius+1), (center[0]-radius, center[1]-radius)) 

85 self.bbox = bbox 

86 self.morph = np.zeros(bbox.shape, dtype=np.float32) 

87 _center = (self.morph.shape[0]-1)//2, (self.morph.shape[1]-1)//2 

88 self.morph[*_center] = 1 

89 

90 

91def initData( 

92 models: list[DeblenderTestModel], 

93 modelPsf: np.ndarray, 

94 imagePsf: np.ndarray, 

95) -> tuple[scl.Image, scl.Image]: 

96 # Find the bounding box that contains all of the models 

97 bbox = models[0].bbox 

98 bands = models[0].bands 

99 for model in models[1:]: 

100 bbox = bbox | model.bbox 

101 assert bands == model.bands 

102 bbox = bbox.grow(5) 

103 

104 for dim in bbox.origin: 

105 if dim < 0: 

106 raise ValueError("Invalid setup, at least one source has a footprint below (0, 0)") 

107 

108 deconvolved = scl.Image.from_box(bbox, bands=bands, dtype=modelPsf.dtype) 

109 convolved = scl.Image.from_box(bbox, bands=bands, dtype=imagePsf.dtype) 

110 for model in models: 

111 deconvolved += model.render(modelPsf) 

112 convolved += model.render(imagePsf) 

113 return deconvolved, convolved