Coverage for python/lsst/scarlet/lite/component.py: 30%

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 "Component", 

26 "FactorizedComponent", 

27 "default_fista_parameterization", 

28 "default_adaprox_parameterization", 

29] 

30 

31from abc import ABC, abstractmethod 

32from functools import partial 

33from typing import TYPE_CHECKING, Callable, cast 

34 

35import numpy as np 

36 

37from .bbox import Box 

38from .image import Image 

39from .operators import Monotonicity, prox_uncentered_symmetry 

40from .parameters import AdaproxParameter, FistaParameter, Parameter, parameter, relative_step 

41 

42if TYPE_CHECKING: 

43 from .io import ScarletComponentBaseData 

44 

45 

46class Component(ABC): 

47 """A base component in scarlet lite. 

48 

49 Parameters 

50 ---------- 

51 bands: 

52 The bands used when the component model is created. 

53 bbox: Box 

54 The bounding box for this component. 

55 """ 

56 

57 def __init__( 

58 self, 

59 bands: tuple, 

60 bbox: Box, 

61 ): 

62 self._bands = bands 

63 self._bbox = bbox 

64 

65 @property 

66 def bbox(self) -> Box: 

67 """The bounding box that contains the component in the full image""" 

68 return self._bbox 

69 

70 @property 

71 def bands(self) -> tuple: 

72 """The bands in the component model""" 

73 return self._bands 

74 

75 @abstractmethod 

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

77 """Test whether or not the component needs to be resized 

78 

79 This should be overriden in inherited classes and return `True` 

80 if the component needs to be resized. 

81 """ 

82 

83 @abstractmethod 

84 def update(self, it: int, input_grad: np.ndarray) -> None: 

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

86 

87 Parameters 

88 ---------- 

89 it: 

90 The current iteration of the optimizer. 

91 input_grad: 

92 Gradient of the likelihood wrt the component model 

93 """ 

94 

95 @abstractmethod 

96 def get_model(self) -> Image: 

97 """Generate a model for the component 

98 

99 This must be implemented in inherited classes. 

100 

101 Returns 

102 ------- 

103 model: Image 

104 The image of the component model. 

105 """ 

106 

107 @abstractmethod 

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

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

110 

111 Parameters 

112 ---------- 

113 parameterization: Callable 

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

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

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

117 """ 

118 

119 @abstractmethod 

120 def to_data(self) -> ScarletComponentBaseData: 

121 """Convert the component to persistable ScarletComponentBaseData 

122 

123 Returns 

124 ------- 

125 component_data: ScarletComponentBaseData 

126 The data object containing the component information 

127 """ 

128 

129 

130class FactorizedComponent(Component): 

131 """A component that can be factorized into spectrum and morphology 

132 parameters. 

133 

134 Parameters 

135 ---------- 

136 bands: 

137 The bands of the spectral dimension, in order. 

138 spectrum: 

139 The parameter to store and update the spectrum. 

140 morph: 

141 The parameter to store and update the morphology. 

142 peak: 

143 Location of the peak for the source. 

144 bbox: 

145 The `Box` in the `model_bbox` that contains the source. 

146 bg_rms: 

147 The RMS of the background used to threshold, grow, 

148 and shrink the component. 

149 bg_thresh: 

150 The threshold to use for the background RMS. 

151 If `None`, no background thresholding is applied, otherwise 

152 a sparsity constraint is applied to the morpholigy that 

153 requires flux in at least one band to be bg_thresh multiplied by 

154 `bg_rms` in that band. 

155 floor: 

156 Minimum value of the spectrum or center morphology pixel 

157 (depending on which is normalized). 

158 monotonicity: 

159 The monotonicity operator to use for making the source monotonic. 

160 If this parameter is `None`, the source will not be made monotonic. 

161 padding: 

162 The amount of padding to add to the component bounding box 

163 when resizing the component. 

164 is_symmetric: 

165 Whether the component is symmetric or not. 

166 If `True`, the morphology will be symmetrized using 

167 `prox_uncentered_symmetry`. 

168 If `False`, the morphology will not be symmetrized. 

169 """ 

170 

171 def __init__( 

172 self, 

173 bands: tuple, 

174 spectrum: Parameter | np.ndarray, 

175 morph: Parameter | np.ndarray, 

176 bbox: Box, 

177 peak: tuple[int, int] | None = None, 

178 bg_rms: np.ndarray | None = None, 

179 bg_thresh: float | None = 0.25, 

180 floor: float = 1e-20, 

181 monotonicity: Monotonicity | None = None, 

182 padding: int = 5, 

183 is_symmetric: bool = False, 

184 ): 

185 # Initialize all of the base attributes 

186 super().__init__( 

187 bands=bands, 

188 bbox=bbox, 

189 ) 

190 self._spectrum = parameter(spectrum) 

191 self._morph = parameter(morph) 

192 self._peak = peak 

193 self.bg_rms = bg_rms 

194 self.bg_thresh = bg_thresh 

195 

196 self.floor = floor 

197 self.monotonicity = monotonicity 

198 self.padding = padding 

199 self.is_symmetric = is_symmetric 

200 

201 @property 

202 def peak(self) -> tuple[int, int] | None: 

203 """The peak of the component 

204 

205 Returns 

206 ------- 

207 peak: 

208 The peak of the component 

209 """ 

210 return self._peak 

211 

212 @property 

213 def component_center(self) -> tuple[int, int] | None: 

214 """The center of the component in its bounding box 

215 

216 This is likely to be different than `Component.center`, 

217 since `Component.center` is the center of the component in the 

218 full model, whereas `component_center` is the center of the component 

219 inside its bounding box. 

220 

221 Returns 

222 ------- 

223 center: 

224 The center of the component in its bounding box 

225 """ 

226 _center = self.peak 

227 if _center is None: 

228 return None 

229 center = ( 

230 _center[0] - self.bbox.origin[-2], 

231 _center[1] - self.bbox.origin[-1], 

232 ) 

233 return center 

234 

235 @property 

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

237 """The array of spectrum values""" 

238 return self._spectrum.x 

239 

240 @property 

241 def morph(self) -> np.ndarray: 

242 """The array of morphology values""" 

243 return self._morph.x 

244 

245 @property 

246 def shape(self) -> tuple: 

247 """Shape of the resulting model image""" 

248 return self.spectrum.shape + self.morph.shape 

249 

250 def get_model(self) -> Image: 

251 """Build the model from the spectrum and morphology""" 

252 # The spectrum and morph might be Parameters, 

253 # so cast them as arrays in the model. 

254 spectrum = self.spectrum 

255 morph = self.morph 

256 model = spectrum[:, None, None] * morph[None, :, :] 

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

258 

259 def grad_spectrum(self, input_grad: np.ndarray, spectrum: np.ndarray, morph: np.ndarray): 

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

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

262 

263 def grad_morph(self, input_grad: np.ndarray, morph: np.ndarray, spectrum: np.ndarray): 

264 """Gradient of the morph wrt. the component model""" 

265 return np.einsum("i,i...", spectrum, input_grad) 

266 

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

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

269 # prevent divergent spectrum 

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

271 spectrum[~np.isfinite(spectrum)] = self.floor 

272 return spectrum 

273 

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

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

276 # Get the peak position in the current bbox 

277 shape = morph.shape 

278 if self.peak is None: 

279 peak = (shape[0] // 2, shape[1] // 2) 

280 else: 

281 peak = ( 

282 self.peak[0] - self.bbox.origin[-2], 

283 self.peak[1] - self.bbox.origin[-1], 

284 ) 

285 

286 # monotonicity 

287 if self.monotonicity is not None: 

288 morph = self.monotonicity(morph, cast(tuple[int, int], self.component_center)) 

289 

290 # symmetry 

291 if self.is_symmetric: 

292 # Apply the symmetry operator 

293 morph = prox_uncentered_symmetry(morph, peak, fill=0.0) 

294 

295 if self.bg_thresh is not None and self.bg_rms is not None: 

296 bg_thresh = self.bg_rms * self.bg_thresh 

297 # Enforce background thresholding 

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

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

300 else: 

301 # enforce positivity 

302 morph[morph < 0] = 0 

303 

304 # prevent divergent morphology 

305 morph[peak] = np.max([morph[peak], self.floor]) 

306 

307 # Ensure that the morphology is finite 

308 morph[~np.isfinite(morph)] = 0 

309 

310 # Normalize the morphology 

311 max_value = np.max(morph) 

312 if max_value > 0: 

313 morph[:] = morph / max_value 

314 return morph 

315 

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

317 """Test whether or not the component needs to be resized""" 

318 # No need to resize if there is no size threshold. 

319 # To allow box sizing but no thresholding use `bg_thresh=0`. 

320 if self.bg_thresh is None or self.bg_rms is None: 

321 return False 

322 

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

324 bg_thresh = self.bg_rms * self.bg_thresh 

325 significant = np.any(model >= bg_thresh[:, None, None], axis=0) 

326 if np.sum(significant) == 0: 

327 # There are no significant pixels, 

328 # so make a small box around the center 

329 center = self.peak 

330 if center is None: 

331 center = (0, 0) 

332 new_box = Box((1, 1), center).grow(self.padding) & model_box 

333 else: 

334 new_box = ( 

335 Box.from_data(significant, threshold=0).grow(self.padding) + self.bbox.origin # type: ignore 

336 ) & model_box 

337 if new_box == self.bbox: 

338 return False 

339 

340 old_box = self.bbox 

341 self._bbox = new_box 

342 self._morph.resize(old_box, new_box) 

343 return True 

344 

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

346 """Update the spectrum and morphology parameters""" 

347 # Store the input spectrum so that the morphology can 

348 # have a consistent update 

349 spectrum = self.spectrum.copy() 

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

351 self._morph.update(it, input_grad, spectrum) 

352 

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

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

355 

356 Parameters 

357 ---------- 

358 parameterization: Callable 

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

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

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

362 """ 

363 # Update the spectrum and morph in place 

364 parameterization(self) 

365 # update the parameters 

366 self._spectrum.grad = self.grad_spectrum 

367 self._spectrum.prox = self.prox_spectrum 

368 self._morph.grad = self.grad_morph 

369 self._morph.prox = self.prox_morph 

370 

371 def to_data(self) -> ScarletComponentBaseData: 

372 """Convert the component to persistable ScarletComponentBaseData 

373 

374 Returns 

375 ------- 

376 component_data: ScarletComponentBaseData 

377 The data object containing the component information 

378 """ 

379 from .io import ScarletFactorizedComponentData 

380 

381 return ScarletFactorizedComponentData( 

382 origin=self.bbox.origin, # type: ignore 

383 peak=self.peak, # type: ignore 

384 spectrum=self.spectrum, 

385 morph=self.morph, 

386 ) 

387 

388 def __str__(self): 

389 result = ( 

390 f"FactorizedComponent<\n bands={self.bands},\n center={self.peak},\n " 

391 f"spectrum={self.spectrum},\n morph_shape={self.morph.shape}\n>" 

392 ) 

393 return result 

394 

395 def __repr__(self): 

396 return self.__str__() 

397 

398 

399def default_fista_parameterization(component: Component): 

400 """Initialize a factorized component to use FISTA PGM for optimization""" 

401 if isinstance(component, FactorizedComponent): 

402 component._spectrum = FistaParameter(component.spectrum, step=0.5) 

403 component._morph = FistaParameter(component.morph, step=0.5) 

404 else: 

405 raise NotImplementedError(f"Unrecognized component type {component}") 

406 

407 

408def default_adaprox_parameterization(component: Component, noise_rms: float | None = None): 

409 """Initialize a factorized component to use Proximal ADAM 

410 for optimization 

411 """ 

412 if noise_rms is None: 

413 noise_rms = 1e-16 

414 if isinstance(component, FactorizedComponent): 

415 component._spectrum = AdaproxParameter( 

416 component.spectrum, 

417 step=partial(relative_step, factor=1e-2, minimum=noise_rms), 

418 ) 

419 component._morph = AdaproxParameter( 

420 component.morph, 

421 step=1e-2, 

422 ) 

423 else: 

424 raise NotImplementedError(f"Unrecognized component type {component}")