Coverage for tests/utils.py: 14%

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 

22import numpy as np 

23import scipy.signal 

24 

25from lsst.geom import Box2I, Point2I, Extent2I 

26from lsst.afw.geom import SpanSet 

27from lsst.afw.detection import Footprint, GaussianPsf 

28import scarlet 

29 

30 

31def numpyToStack(images, center, offset): 

32 """Convert numpy and python objects to stack objects 

33 """ 

34 cy, cx = center 

35 bands, height, width = images.shape 

36 x0, y0 = offset 

37 bbox = Box2I(Point2I(x0, y0), Extent2I(width, height)) 

38 spanset = SpanSet(bbox) 

39 foot = Footprint(spanset) 

40 foot.addPeak(cx+x0, cy+y0, images[:, cy, cx].max()) 

41 peak = foot.getPeaks()[0] 

42 return foot, peak, bbox 

43 

44 

45def initData(shape, coords, amplitudes=None, convolve=True): 

46 """Initialize data for the tests 

47 """ 

48 

49 B, Ny, Nx = shape 

50 K = len(coords) 

51 

52 if amplitudes is None: 

53 amplitudes = np.ones((K,)) 

54 assert K == len(amplitudes) 

55 

56 _seds = [ 

57 np.arange(B, dtype=float), 

58 np.arange(B, dtype=float)[::-1], 

59 np.ones((B,), dtype=float) 

60 ] 

61 seds = np.array([_seds[n % 3]*amplitudes[n] for n in range(K)], dtype=np.float32) 

62 

63 morphs = np.zeros((K, Ny, Nx), dtype=np.float32) 

64 for k, coord in enumerate(coords): 

65 morphs[k, coord[0], coord[1]] = 1 

66 images = seds.T.dot(morphs.reshape(K, -1)).reshape(shape) 

67 

68 if convolve: 

69 psfRadius = 20 

70 psfShape = (2*psfRadius+1, 2*psfRadius+1) 

71 

72 targetPsf = scarlet.GaussianPSF(sigma=0.9, boxsize=2*psfShape[0]) 

73 targetPsfImage = targetPsf.get_model()[0] 

74 

75 psfs = [GaussianPsf(psfShape[1], psfShape[0], 1+.2*b) for b in range(B)] 

76 psfImages = np.array([psf.computeImage().array for psf in psfs]) 

77 psfImages /= psfImages.max(axis=(1, 2))[:, None, None] 

78 

79 # Convolve the image with the psf in each channel 

80 # Use scipy.signal.convolve without using FFTs as a sanity check 

81 images = np.array([scipy.signal.convolve(img, psf, method="direct", mode="same") 

82 for img, psf in zip(images, psfImages)]) 

83 # Convolve the true morphology with the target PSF, 

84 # also using scipy.signal.convolve as a sanity check 

85 morphs = np.array([scipy.signal.convolve(m, targetPsfImage, method="direct", mode="same") 

86 for m in morphs]) 

87 morphs /= morphs.max() 

88 psfImages /= psfImages.sum(axis=(1, 2))[:, None, None] 

89 psfImages = scarlet.ImagePSF(psfImages) 

90 

91 channels = range(len(images)) 

92 return targetPsfImage, psfImages, images, channels, seds, morphs, targetPsf, psfs