Coverage for python/lsst/meas/extensions/scarlet/footprint.py: 98%

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 

22"""Functions for converting between afw and scarlet footprints.""" 

23 

24from typing import Sequence 

25 

26import lsst.geom as geom 

27import lsst.scarlet.lite as scl 

28import numpy as np 

29from lsst.afw.detection import Footprint as afwFootprint 

30from lsst.afw.detection import HeavyFootprintF, PeakCatalog, makeHeavyFootprint 

31from lsst.afw.detection.multiband import MultibandFootprint 

32from lsst.afw.geom import SpanSet 

33from lsst.afw.image import Image as afwImage 

34from lsst.afw.image import Mask, MaskedImage, MultibandImage 

35from lsst.scarlet.lite.detect_pybind11 import Peak 

36 

37from .utils import bboxToScarletBox, nonzeroBandSupport 

38 

39 

40__all__ = [ 

41 "afwFootprintToScarlet", 

42 "scarletFootprintToAfw", 

43 "scarletModelToHeavy", 

44 "scarletFootprintsToPeakCatalog", 

45] 

46 

47 

48def afwFootprintToScarlet(footprint: afwFootprint, copyPeaks: bool = True): 

49 """Convert an afw Footprint into a scarlet lite Footprint. 

50 

51 Parameters 

52 ---------- 

53 footprint: 

54 The afw Footprint to convert. 

55 copyPeaks: 

56 Whether or not to copy the peaks from the afw Footprint. 

57 

58 Returns 

59 ------- 

60 scarletFootprint: 

61 The converted scarlet Footprint. 

62 """ 

63 data = footprint.spans.asArray() 

64 afwBox = footprint.getBBox() 

65 bbox = bboxToScarletBox(afwBox) 

66 peaks = [] 

67 if copyPeaks: 67 ↛ 71line 67 didn't jump to line 71 because the condition on line 67 was always true

68 for peak in footprint.peaks: 

69 newPeak = Peak(peak.getIy(), peak.getIx(), peak.getPeakValue()) 

70 peaks.append(newPeak) 

71 bounds = scl.detect.bbox_to_bounds(bbox) 

72 return scl.detect.Footprint(data, peaks, bounds) 

73 

74 

75def scarletFootprintToAfw(footprint: scl.detect.Footprint, copyPeaks: bool = True) -> afwFootprint: 

76 """Convert a scarlet lite Footprint into an afw Footprint. 

77 

78 Parameters 

79 ---------- 

80 footprint: 

81 The scarlet Footprint to convert. 

82 copyPeaks: 

83 Whether or not to copy the peaks from the scarlet Footprint. 

84 

85 Returns 

86 ------- 

87 newFootprint: 

88 The converted afw Footprint. 

89 """ 

90 xy0 = geom.Point2I(footprint.bbox.origin[1], footprint.bbox.origin[0]) 

91 data = Mask(footprint.data.astype(np.int32), xy0=xy0) 

92 spans = SpanSet.fromMask(data) 

93 newFootprint = afwFootprint(spans) 

94 

95 if copyPeaks: 95 ↛ 98line 95 didn't jump to line 98 because the condition on line 95 was always true

96 for peak in footprint.peaks: 

97 newFootprint.addPeak(peak.x, peak.y, peak.flux) 

98 return newFootprint 

99 

100 

101def scarletModelToHeavy( 

102 source: scl.Source, 

103 blend: scl.Blend, 

104 useFlux=False, 

105) -> HeavyFootprintF | MultibandFootprint: 

106 """Convert a scarlet_lite model to a `HeavyFootprintF` 

107 or `MultibandFootprint`. 

108 

109 Parameters 

110 ---------- 

111 source: 

112 The source to convert to a `HeavyFootprint`. 

113 blend: 

114 The `Blend` object that contains information about 

115 the observation, PSF, etc, used to convolve the 

116 scarlet model to the observed seeing in each band. 

117 useFlux: 

118 Whether or not to re-distribute the flux from the image 

119 to conserve flux. 

120 

121 Returns 

122 ------- 

123 heavy: 

124 The footprint (possibly multiband) containing the model for the source. 

125 """ 

126 # We want to convolve the model with the observed PSF, 

127 # which means we need to grow the model box by the PSF to 

128 # account for all of the flux after convolution. 

129 

130 # Get the PSF size and radii to grow the box 

131 py, px = blend.observation.psfs.shape[1:] 

132 dh = py // 2 

133 dw = px // 2 

134 

135 if useFlux: 

136 bbox = source.flux_weighted_image.bbox 

137 else: 

138 bbox = source.bbox.grow((dh, dw)) 

139 # Only use the portion of the convolved model that fits in the image 

140 overlap = bbox & blend.observation.bbox 

141 # Load the full multiband model in the larger box 

142 if useFlux: 

143 # The flux weighted model is already convolved, so we just load it 

144 model = source.get_model(use_flux=True).project(bbox=overlap) 

145 else: 

146 model = source.get_model().project(bbox=overlap) 

147 # Convolve the model with the PSF in each band 

148 # Always use a real space convolution to limit artifacts 

149 model = blend.observation.convolve(model, mode="real") 

150 

151 # Update xy0 with the origin of the sources box 

152 xy0 = geom.Point2I(model.yx0[-1], model.yx0[-2]) 

153 # Create the spans for the footprint 

154 valid = nonzeroBandSupport(model.data) 

155 valid = Mask(valid.astype(np.int32), xy0=xy0) 

156 spans = SpanSet.fromMask(valid) 

157 

158 # Create the MultibandHeavyFootprint and 

159 # add the location of the source to the peak catalog. 

160 foot = afwFootprint(spans) 

161 foot.addPeak(source.center[1], source.center[0], np.max(model.data)) 

162 if model.n_bands == 1: 

163 image = afwImage( 

164 array=model.data[0], xy0=valid.getBBox().getMin(), dtype=model.dtype 

165 ) 

166 maskedImage = MaskedImage(image, dtype=model.dtype) 

167 heavy = makeHeavyFootprint(foot, maskedImage) 

168 else: 

169 bands = model.bands 

170 model = MultibandImage(bands, model.data, valid.getBBox()) 

171 heavy = MultibandFootprint.fromImages(bands, model, footprint=foot) 

172 return heavy 

173 

174 

175def scarletFootprintsToPeakCatalog( 

176 footprints: Sequence[scl.detect.Footprint], 

177) -> PeakCatalog: 

178 """Create a PeakCatalog from a list of scarlet footprints. 

179 

180 This creates a dummy Footprint to add the peaks to, 

181 then extracts the peaks from the Footprint. 

182 It would be better to create a PeakCatalog directly, 

183 but currently that is not supported in afw. 

184 

185 Parameters 

186 ---------- 

187 footprints: 

188 A list of scarlet footprints. 

189 

190 Returns 

191 ------- 

192 peaks: 

193 A PeakCatalog containing all of the peaks in the footprints. 

194 """ 

195 tempFootprint = afwFootprint() 

196 for footprint in footprints: 

197 for peak in footprint.peaks: 

198 tempFootprint.addPeak(peak.x, peak.y, peak.flux) 

199 return tempFootprint.peaks