Coverage for python/lsst/scarlet/lite/models/free_form.py: 27%

94 statements  

« prev     ^ index     » next       coverage.py v7.11.0, created at 2025-11-07 01:40 -0800

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__ = ["FactorizedFreeFormComponent"] 

23 

24from typing import Callable, cast 

25 

26import numpy as np 

27from lsst.scarlet.lite.detect_pybind11 import get_connected_multipeak, get_footprints # type: ignore 

28 

29from ..bbox import Box 

30from ..component import Component, FactorizedComponent 

31from ..detect import footprints_to_image 

32from ..image import Image 

33from ..parameters import Parameter, parameter 

34 

35 

36class FactorizedFreeFormComponent(FactorizedComponent): 

37 """Implements a free-form component 

38 

39 With no constraints this component is typically either a garbage collector, 

40 or part of a set of components to deconvolve an image by separating out 

41 the different spectral components. 

42 

43 See `FactorizedComponent` for a list of parameters not shown here. 

44 

45 Parameters 

46 ---------- 

47 peaks: `list` of `tuple` 

48 A set of ``(cy, cx)`` peaks for detected sources. 

49 If peak is not ``None`` then only pixels in the same "footprint" 

50 as one of the peaks are included in the morphology. 

51 If `peaks` is ``None`` then there is no constraint applied. 

52 min_area: float 

53 The minimum area for a peak. 

54 If `min_area` is not `None` then all regions of the morphology 

55 with fewer than `min_area` connected pixels are removed. 

56 """ 

57 

58 def __init__( 

59 self, 

60 bands: tuple, 

61 spectrum: np.ndarray | Parameter, 

62 morph: np.ndarray | Parameter, 

63 model_bbox: Box, 

64 bg_thresh: float | None = None, 

65 bg_rms: np.ndarray | None = None, 

66 floor: float = 1e-20, 

67 peaks: list[tuple[int, int]] | None = None, 

68 min_area: float = 0, 

69 ): 

70 super().__init__( 

71 bands=bands, 

72 spectrum=spectrum, 

73 morph=morph, 

74 bbox=model_bbox, 

75 peak=None, 

76 bg_rms=bg_rms, 

77 bg_thresh=bg_thresh, 

78 floor=floor, 

79 ) 

80 

81 self.peaks = peaks 

82 self.min_area = min_area 

83 

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

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

86 

87 This differs from `FactorizedComponent` because an 

88 `SedComponent` has the spectrum normalized to unity. 

89 """ 

90 # prevent divergent spectrum 

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

92 # Normalize the spectrum 

93 spectrum = spectrum / np.sum(spectrum) 

94 return spectrum 

95 

96 def prox_morph(self, morph: np.ndarray) -> np.ndarray: 

97 """Apply a prox-like update to the morphology 

98 

99 This is the main difference between an `SedComponent` and a 

100 `FactorizedComponent`, since this component has fewer constraints. 

101 """ 

102 from lsst.scarlet.lite.detect_pybind11 import get_connected_multipeak, get_footprints # type: ignore 

103 

104 if self.bg_thresh is not None and isinstance(self.bg_rms, np.ndarray): 

105 bg_thresh = self.bg_rms * self.bg_thresh 

106 # Enforce background thresholding 

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

108 morph[np.all(model < bg_thresh[:, None, None], axis=0)] = 0 

109 else: 

110 # enforce positivity 

111 morph[morph < 0] = 0 

112 

113 if self.peaks is not None: 

114 footprint = get_connected_multipeak(morph, self.peaks, 0) 

115 morph = morph * footprint 

116 

117 if self.min_area > 0: 

118 footprints = get_footprints(morph, 4.0, self.min_area, 0, 0, False) 

119 bbox = self.bbox.copy() 

120 bbox.origin = (0, 0) 

121 footprint_image = footprints_to_image(footprints, bbox) 

122 morph = morph * (footprint_image > 0).data 

123 

124 if np.all(morph == 0): 

125 morph[0, 0] = self.floor 

126 

127 return morph 

128 

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

130 return False 

131 

132 def __str__(self): 

133 return ( 

134 f"FactorizedFreeFormComponent(\n bands={self.bands}\n " 

135 f"spectrum={self.spectrum})\n center={self.peak}\n " 

136 f"morph_shape={self.morph.shape}" 

137 ) 

138 

139 def __repr__(self): 

140 return self.__str__() 

141 

142 

143class FreeFormComponent(Component): 

144 """Implements a component with no spectral or monotonicty constraints 

145 

146 This is a FreeFormComponent that is not factorized into a 

147 spectrum and morphology with no monotonicity constraint. 

148 """ 

149 

150 def __init__( 

151 self, 

152 bands: tuple, 

153 model: np.ndarray | Parameter, 

154 model_bbox: Box, 

155 bg_thresh: float | None = None, 

156 bg_rms: np.ndarray | None = None, 

157 floor: float = 1e-20, 

158 peaks: list[tuple[int, int]] | None = None, 

159 min_area: float = 0, 

160 ): 

161 super().__init__(bands=bands, bbox=model_bbox) 

162 self._model = parameter(model) 

163 self.bg_rms = bg_rms 

164 self.bg_thresh = bg_thresh 

165 self.floor = floor 

166 self.peaks = peaks 

167 self.min_area = min_area 

168 

169 @property 

170 def model(self) -> np.ndarray: 

171 return self._model.x 

172 

173 def get_model(self) -> Image: 

174 return Image(self.model, bands=self.bands, yx0=cast(tuple[int, int], self.bbox.origin)) 

175 

176 @property 

177 def shape(self) -> tuple: 

178 return self.model.shape 

179 

180 def grad_model(self, input_grad: np.ndarray, model: np.ndarray) -> np.ndarray: 

181 return input_grad 

182 

183 def prox_model(self, model: np.ndarray) -> np.ndarray: 

184 if self.bg_thresh is not None and isinstance(self.bg_rms, np.ndarray): 

185 bg_thresh = self.bg_rms * self.bg_thresh 

186 # Enforce background thresholding 

187 model[model < bg_thresh[:, None, None]] = 0 

188 else: 

189 # enforce positivity 

190 model[model < 0] = 0 

191 

192 if self.peaks is not None: 

193 # Remove pixels not connected to one of the peaks 

194 model2d = np.sum(model, axis=0) 

195 footprint = get_connected_multipeak(model2d, self.peaks, 0) 

196 model = model * footprint[None, :, :] 

197 

198 if self.min_area > 0: 

199 # Remove regions with fewer than min_area connected pixels 

200 model2d = np.sum(model, axis=0) 

201 footprints = get_footprints(model2d, 4.0, self.min_area, 0, 0, False) 

202 bbox = self.bbox.copy() 

203 bbox.origin = (0, 0) 

204 footprint_image = footprints_to_image(footprints, bbox) 

205 model = model * (footprint_image > 0).data[None, :, :] 

206 

207 if np.all(model == 0): 

208 # If the model is all zeros, set a single pixel to the floor 

209 model[0, 0] = self.floor 

210 

211 return model 

212 

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

214 return False 

215 

216 def update(self, it: int, grad_log_likelihood: np.ndarray): 

217 self._model.update(it, grad_log_likelihood) 

218 

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

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

221 

222 Parameters 

223 ---------- 

224 parameterization: Callable 

225 A function to use to convert parameters of a given type into 

226 a `Parameter` in place. It should take a single argument that 

227 is the `Component` or `Source` that is to be parameterized. 

228 """ 

229 # Update the spectrum and morph in place 

230 parameterization(self) 

231 # update the parameters 

232 self._model.grad = self.grad_model 

233 self._model.prox = self.prox_model 

234 

235 def __str__(self): 

236 result = f"FreeFormComponent<bands={self.bands}, shape={self.shape}>" 

237 return result 

238 

239 def __repr__(self): 

240 return self.__str__()