Coverage for python / lsst / scarlet / lite / source.py: 37%

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

25 

26from abc import ABC, abstractmethod 

27from copy import deepcopy 

28from typing import TYPE_CHECKING, Any, Callable, Self 

29 

30from .bbox import Box 

31from .component import Component 

32from .image import Image 

33 

34if TYPE_CHECKING: 

35 from .io import ScarletSourceBaseData, ScarletSourceData 

36 

37 

38class SourceBase(ABC): 

39 """Base class for a source 

40 

41 This is primarily to allow `isinstance` checks 

42 without importing the full `Source` class. 

43 """ 

44 

45 metadata: dict[str, Any] | None = None 

46 components: list[Component] 

47 

48 @abstractmethod 

49 def to_data(self) -> ScarletSourceBaseData: 

50 """Convert to a `ScarletSourceBaseData` for serialization 

51 

52 Returns 

53 ------- 

54 source_data: 

55 The `ScarletSourceData` representation of this source. 

56 """ 

57 

58 @abstractmethod 

59 def __getitem__(self, indices: Any) -> Self: 

60 """Get a sub-source corresponding to the given indices. 

61 

62 Parameters 

63 ---------- 

64 indices: Any 

65 The indices to use to slice the source model. 

66 

67 Returns 

68 ------- 

69 source: SourceBase 

70 A new source that is a sub-source of this one. 

71 

72 Raises 

73 ------ 

74 IndexError : 

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

76 """ 

77 

78 @abstractmethod 

79 def __deepcopy__(self, memo: dict[int, Any]) -> Self: 

80 """Create a deep copy of this source. 

81 

82 Parameters 

83 ---------- 

84 memo : dict[int, Any] 

85 A memoization dictionary used by `copy.deepcopy`. 

86 

87 Returns 

88 ------- 

89 source : SourceBase 

90 A new source that is a deep copy of this one. 

91 """ 

92 

93 @abstractmethod 

94 def __copy__(self) -> Self: 

95 """Create a copy of this source. 

96 

97 Returns 

98 ------- 

99 source : SourceBase 

100 A new source that is a copy of this one. 

101 """ 

102 

103 def copy(self, deep: bool = False) -> Self: 

104 """Create a copy of this source. 

105 

106 Parameters 

107 ---------- 

108 deep : bool, optional 

109 If `True`, a deep copy is made. If `False`, a shallow copy is made. 

110 Default is `False`. 

111 

112 Returns 

113 ------- 

114 source : Self 

115 A new source that is a copy of this one. 

116 """ 

117 if deep: 

118 return self.__deepcopy__({}) 

119 return self.__copy__() 

120 

121 

122class Source(SourceBase): 

123 """A container for components associated with the same astrophysical object 

124 

125 A source can have a single component, or multiple components, 

126 and each can be contained in different bounding boxes. 

127 

128 Parameters 

129 ---------- 

130 components: 

131 The components contained in the source. 

132 """ 

133 

134 def __init__( 

135 self, 

136 components: list[Component], 

137 metadata: dict | None = None, 

138 flux_weighted_image: Image | None = None, 

139 ): 

140 self.components = components 

141 self.flux_weighted_image = flux_weighted_image 

142 self.metadata = metadata 

143 

144 @property 

145 def n_components(self) -> int: 

146 """The number of components in this source""" 

147 return len(self.components) 

148 

149 @property 

150 def center(self) -> tuple[int, int] | None: 

151 """The center of the source in the full Blend.""" 

152 if not self.is_null and hasattr(self.components[0], "peak"): 

153 return self.components[0].peak # type: ignore 

154 return None 

155 

156 @property 

157 def source_center(self) -> tuple[int, int] | None: 

158 """The center of the source in its local bounding box.""" 

159 _center = self.center 

160 _origin = self.bbox.origin 

161 if _center is not None: 

162 center = ( 

163 _center[0] - _origin[0], 

164 _center[1] - _origin[1], 

165 ) 

166 return center 

167 return None 

168 

169 @property 

170 def is_null(self) -> bool: 

171 """True if the source does not have any components""" 

172 return self.n_components == 0 

173 

174 @property 

175 def bbox(self) -> Box: 

176 """The minimal bounding box to contain all of this sources components 

177 

178 Null sources have a bounding box with shape `(0,0,0)` 

179 """ 

180 if self.n_components == 0: 

181 return Box((0, 0)) 

182 bbox = self.components[0].bbox 

183 for component in self.components[1:]: 

184 bbox = bbox | component.bbox 

185 return bbox 

186 

187 @property 

188 def bands(self) -> tuple: 

189 """The bands in the full source model.""" 

190 if self.is_null: 

191 return () 

192 return self.components[0].bands 

193 

194 def get_model(self, use_flux: bool = False) -> Image: 

195 """Build the model for the source 

196 

197 This is never called during optimization and is only used 

198 to generate a model of the source for investigative purposes. 

199 

200 Parameters 

201 ---------- 

202 use_flux: 

203 Whether to use the re-distributed flux associated with the source 

204 instead of the component models. 

205 

206 Returns 

207 ------- 

208 model: 

209 The full-color model. 

210 """ 

211 if self.n_components == 0: 

212 return 0 # type: ignore 

213 

214 if use_flux: 

215 # Return the redistributed flux 

216 # (calculated by scarlet.lite.measure.weight_sources) 

217 return self.flux_weighted_image # type: ignore 

218 

219 model = self.components[0].get_model() 

220 for component in self.components[1:]: 

221 model = model + component.get_model() 

222 return model 

223 

224 def parameterize(self, parameterization: Callable): 

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

226 

227 Parameters 

228 ---------- 

229 parameterization: 

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

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

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

233 """ 

234 for component in self.components: 

235 component.parameterize(parameterization) 

236 

237 def to_data(self) -> ScarletSourceData: 

238 """Convert to a `ScarletSourceData` for serialization 

239 

240 Returns 

241 ------- 

242 source_data: 

243 The `ScarletSourceData` representation of this source. 

244 """ 

245 from .io import ScarletSourceData 

246 

247 component_data = [c.to_data() for c in self.components] 

248 return ScarletSourceData(components=component_data, metadata=self.metadata) 

249 

250 def __str__(self): 

251 return f"Source<{len(self.components)}>" 

252 

253 def __repr__(self): 

254 return f"Source(components={repr(self.components)})>" 

255 

256 def __getitem__(self, indices: Any) -> Source: 

257 """Get a sub-source corresponding to the given indices. 

258 

259 Parameters 

260 ---------- 

261 indices: Any 

262 The indices to use to slice the source model. Can be: 

263 - A single band 

264 - A slice with start/stop bands 

265 - A sequence of bands 

266 

267 Returns 

268 ------- 

269 source: Source 

270 A new source that is a sub-source of this one. 

271 

272 Raises 

273 ------ 

274 IndexError : 

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

276 """ 

277 flux = None if self.flux_weighted_image is None else self.flux_weighted_image[indices] 

278 return Source( 

279 components=[c[indices] for c in self.components], 

280 metadata=self.metadata, 

281 flux_weighted_image=flux, 

282 ) 

283 

284 def __deepcopy__(self, memo: dict[int, Any]) -> Source: 

285 """Create a deep copy of this source. 

286 

287 Parameters 

288 ---------- 

289 memo : dict[int, Any] 

290 A memoization dictionary used by `copy.deepcopy`. 

291 

292 Returns 

293 ------- 

294 source : SourceBase 

295 A new source that is a deep copy of this one. 

296 """ 

297 # Check if already copied 

298 if id(self) in memo: 

299 return memo[id(self)] 

300 

301 # Create placeholder and add to memo FIRST 

302 source = Source.__new__(Source) 

303 memo[id(self)] = source 

304 

305 source.__init__( # type: ignore[misc] 

306 components=deepcopy(self.components, memo), 

307 metadata=deepcopy(self.metadata, memo), 

308 flux_weighted_image=deepcopy(self.flux_weighted_image, memo), 

309 ) 

310 return source 

311 

312 def __copy__(self) -> Source: 

313 """Create a copy of this source. 

314 

315 Returns 

316 ------- 

317 source : SourceBase 

318 A new source that is a copy of this one. 

319 """ 

320 source = Source( 

321 components=self.components, 

322 metadata=self.metadata, 

323 flux_weighted_image=self.flux_weighted_image, 

324 ) 

325 return source