Coverage for src/ufig/psf_estimation/cutouts_utils.py: 100%

1# Copyright (c) 2017 ETH Zurich, Cosmology Research Group 

2""" 

3Created on Feb 23, 2021 

4@author: Tomasz Kacprzak 

5""" 

6 

7import numba as nb 

8import numpy as np 

9 

10 

11@nb.jit(nopython=True) 

12def find_max_flux_pos(x, y, delta_x, delta_y, image): 

13 """ 

14 Find the position of the maximum flux in a 2D image around a given position (x, y). 

15 

16 :param x: x-coordinate of the center position 

17 :param y: y-coordinate of the center position 

18 :param delta_x: array of x offsets to check 

19 :param delta_y: array of y offsets to check 

20 :param image: 2D numpy array representing the image 

21 :return: (xi_max, yi_max) - coordinates of the pixel with maximum flux 

22 """ 

23 

24 max_flux = -np.inf 

25 xi_max = 0 

26 yi_max = 0 

27 

28 for dx in delta_x: 

29 for dy in delta_y: 

30 xi = x + dx 

31 yi = y + dy 

32 

33 if ( 

34 0 <= xi < image.shape[1] 

35 and 0 <= yi < image.shape[0] 

36 and image[yi, xi] > max_flux 

37 ): 

38 max_flux = image[yi, xi] 

39 xi_max = xi 

40 yi_max = yi 

41 

42 return xi_max, yi_max 

43 

44 

45def get_cutouts(x_peak, y_peak, image, pointings_maps, stamp_shape): 

46 """ 

47 Get cube of coutouts around a list of coordinates. 

48 

49 :param x_peak: list of x positions 

50 :param y_peak: list of y positions 

51 :param image: 2D image 

52 :param pointings_maps: an array with binary pointings indicators 

53 :param stamp_shape: (nx, ny) shape of stamp, fixed for all objects 

54 """ 

55 

56 stamp_half_x = stamp_shape[1] // 2 

57 stamp_half_y = stamp_shape[0] // 2 

58 

59 delta_x_check = np.arange(-1, 2, dtype=np.int32) 

60 delta_y_check = np.arange(-1, 2, dtype=np.int32) 

61 

62 cube = np.zeros((len(x_peak),) + stamp_shape, dtype=np.float32) 

63 

64 select_image_boundary = np.zeros(len(x_peak), dtype=bool) 

65 select_coadd_boundary = np.zeros(len(x_peak), dtype=bool) 

66 select_max_flux_centered = np.zeros(len(x_peak), dtype=bool) 

67 

68 for i_obj in range(len(x_peak)): 

69 # 1) select pixel with maximum intensity 

70 xi_max, yi_max = find_max_flux_pos( 

71 x_peak[i_obj], y_peak[i_obj], delta_x_check, delta_y_check, image 

72 ) 

73 

74 # 2) compute indices for cutting out 

75 x_start = xi_max - stamp_half_x 

76 y_start = yi_max - stamp_half_y 

77 x_stop = xi_max + stamp_half_x + 1 

78 y_stop = yi_max + stamp_half_y + 1 

79 

80 # 3) check if star is completely on image 

81 select_image_boundary[i_obj] = ( 

82 (x_start >= 0) 

83 & (y_start >= 0) 

84 & (x_stop < image.shape[1]) 

85 & (y_stop < image.shape[0]) 

86 ) 

87 

88 # 4) check if star hits any chip boundaries 

89 pointings = pointings_maps[y_start:y_stop, x_start:x_stop] 

90 select_coadd_boundary[i_obj] = len(np.unique(pointings)) == 1 

91 

92 # 5) check if cutout is centered on maximum flux (if full stamp is in image) 

93 if select_image_boundary[i_obj]: 

94 cutout = image[y_start:y_stop, x_start:x_stop] 

95 select_max_flux_centered[i_obj] = cutout[ 

96 stamp_half_y, stamp_half_x 

97 ] == np.amax(cutout) 

98 else: 

99 select_max_flux_centered[i_obj] = False 

100 if select_max_flux_centered[i_obj]: 

101 cube[i_obj, :, :] = cutout 

102 

103 return cube, select_image_boundary, select_coadd_boundary, select_max_flux_centered