Coverage for tests/utils.py: 14%

43 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-03-22 03:02 -0700

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 lsst.scarlet.lite as scl 

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 x = np.arange(-psfRadius, psfRadius, psfShape[0]) 

72 y = x.copy() 

73 

74 targetPsfImage = scl.utils.integrated_circular_gaussian(sigma=0.9, x=x, y=y) 

75 

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

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

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

79 

80 # Convolve the image with the psf in each channel 

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

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

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

84 # Convolve the true morphology with the target PSF, 

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

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

87 for m in morphs]) 

88 morphs /= morphs.max() 

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

90 

91 channels = range(len(images)) 

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