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

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 

38 

39 

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

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

42 

43 Parameters 

44 ---------- 

45 footprint: 

46 The afw Footprint to convert. 

47 copyPeaks: 

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

49 

50 Returns 

51 ------- 

52 scarletFootprint: 

53 The converted scarlet Footprint. 

54 """ 

55 data = footprint.spans.asArray() 

56 afwBox = footprint.getBBox() 

57 bbox = bboxToScarletBox(afwBox) 

58 peaks = [] 

59 if copyPeaks: 

60 for peak in footprint.peaks: 

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

62 peaks.append(newPeak) 

63 bounds = scl.detect.bbox_to_bounds(bbox) 

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

65 

66 

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

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

69 

70 Parameters 

71 ---------- 

72 footprint: 

73 The scarlet Footprint to convert. 

74 copyPeaks: 

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

76 

77 Returns 

78 ------- 

79 newFootprint: 

80 The converted afw Footprint. 

81 """ 

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

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

84 spans = SpanSet.fromMask(data) 

85 newFootprint = afwFootprint(spans) 

86 

87 if copyPeaks: 

88 for peak in footprint.peaks: 

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

90 return newFootprint 

91 

92 

93def scarletModelToHeavy( 

94 source: scl.Source, 

95 blend: scl.Blend, 

96 useFlux=False, 

97) -> HeavyFootprintF | MultibandFootprint: 

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

99 or `MultibandFootprint`. 

100 

101 Parameters 

102 ---------- 

103 source: 

104 The source to convert to a `HeavyFootprint`. 

105 blend: 

106 The `Blend` object that contains information about 

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

108 scarlet model to the observed seeing in each band. 

109 useFlux: 

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

111 to conserve flux. 

112 

113 Returns 

114 ------- 

115 heavy: 

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

117 """ 

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

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

120 # account for all of the flux after convolution. 

121 

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

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

124 dh = py // 2 

125 dw = px // 2 

126 

127 if useFlux: 

128 bbox = source.flux_weighted_image.bbox 

129 else: 

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

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

132 overlap = bbox & blend.observation.bbox 

133 # Load the full multiband model in the larger box 

134 if useFlux: 

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

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

137 else: 

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

139 # Convolve the model with the PSF in each band 

140 # Always use a real space convolution to limit artifacts 

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

142 

143 # Update xy0 with the origin of the sources box 

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

145 # Create the spans for the footprint 

146 valid = np.max(model.data, axis=0) != 0 

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

148 spans = SpanSet.fromMask(valid) 

149 

150 # Create the MultibandHeavyFootprint and 

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

152 foot = afwFootprint(spans) 

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

154 if model.n_bands == 1: 

155 image = afwImage( 

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

157 ) 

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

159 heavy = makeHeavyFootprint(foot, maskedImage) 

160 else: 

161 model = MultibandImage(blend.bands, model.data, valid.getBBox()) 

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

163 return heavy 

164 

165 

166def scarletFootprintsToPeakCatalog( 

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

168) -> PeakCatalog: 

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

170 

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

172 then extracts the peaks from the Footprint. 

173 It would be better to create a PeakCatalog directly, 

174 but currently that is not supported in afw. 

175 

176 Parameters 

177 ---------- 

178 footprints: 

179 A list of scarlet footprints. 

180 

181 Returns 

182 ------- 

183 peaks: 

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

185 """ 

186 tempFootprint = afwFootprint() 

187 for footprint in footprints: 

188 for peak in footprint.peaks: 

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

190 return tempFootprint.peaks