Coverage for src/ufig/plugins/background_noise.py: 100%

1# Copyright (c) 2013 ETH Zurich, Institute of Astronomy, Lukas Gamper 

2# <lukas.gamper@usystems.ch> 

3""" 

4Created on Oct 7, 2013 

5@author: Lukas Gamper 

6""" 

7 

8import h5py 

9import numpy as np 

10from cosmic_toolbox import logger 

11from ivy.plugin.base_plugin import BasePlugin 

12 

13from ufig import io_util, sysmaps_util 

14 

15LOGGER = logger.get_logger(__file__) 

16 

17 

18def get_effective_bkg_noise_scale_factor(par): 

19 bkg_noise_factor = par.gain if par.bkg_noise_multiply_gain else 1 

20 return bkg_noise_factor 

21 

22 

23def add_from_gaussian(ctx): 

24 """ 

25 Adds noise to the image assuming a Gaussian background model. 

26 The noise in every pixel is randomly drawn from a Gaussian 

27 

28 :param ctx.image: Simulated image 

29 :param ctx.parameters.bkg_noise_amp: Center of the Gaussian distribution 

30 :param ctx.parameters.background_sigma: RMS of the Gaussian distribution 

31 """ 

32 par = ctx.parameters 

33 bkg_noise_factor = get_effective_bkg_noise_scale_factor(ctx.parameters) 

34 amp_variations = np.random.normal( 

35 0, par.bkg_amp_variation_sigma, size=ctx.image.shape 

36 ) 

37 noise_variations = np.random.normal( 

38 0, par.bkg_noise_variation_sigma, size=ctx.image.shape 

39 ) 

40 par.bkg_noise_std = bkg_noise_factor * par.background_sigma + noise_variations 

41 select = par.bkg_noise_std < 0 

42 if np.any(select): 

43 LOGGER.warning("Negative noise std values detected, setting to 0") 

44 par.bkg_noise_std[select] = 0 

45 ctx.image += np.random.normal( 

46 par.bkg_noise_amp + amp_variations, 

47 par.bkg_noise_std, 

48 size=ctx.image.shape, 

49 ) 

50 ctx.image_mask = np.ones((par.size_y, par.size_x), dtype=bool) 

51 

52 

53def add_from_map(ctx): 

54 """ 

55 Adds noise to the image using a rms map 

56 

57 :param ctx.image: Simulated image 

58 :param ctx.params.maps_remote_dir: Path to fits-image containing the RMS of the 

59 noise in every pixel of the image 

60 :param ctx.params.bkg_rms_file_name: Path to fits-image containing the RMS of the 

61 noise in every pixel of the image :param 

62 ctx.params.bkg_noise_scale: Scale factor 

63 applied to the map 

64 :param ctx.params.bkg_noise_amp: Amplitude factor applied to the map 

65 """ 

66 

67 par = ctx.parameters 

68 img_shape = ctx.image.shape 

69 

70 filename_h5, dataset = sysmaps_util.get_hdf_location_bgrms(par) 

71 filepath_h5 = io_util.get_abs_path(filename_h5, par.maps_remote_dir) 

72 

73 bkg_rms = io_util.load_from_hdf5( 

74 file_name=filepath_h5, 

75 hdf5_keys=dataset, 

76 hdf5_path="", 

77 root_path=par.maps_remote_dir, 

78 ) 

79 

80 # Variable systematics 

81 variable_amp = np.random.normal(0, par.bkg_amp_variation_sigma, size=img_shape) 

82 variable_noise = np.random.normal(0, par.bkg_noise_variation_sigma, size=img_shape) 

83 

84 bkg_noise_factor = get_effective_bkg_noise_scale_factor(par) 

85 noise_draw = np.random.normal(0, 1.0, size=img_shape) 

86 par.bkg_noise_std = ( 

87 bkg_rms * bkg_noise_factor * par.bkg_noise_scale + variable_noise 

88 ) 

89 select = par.bkg_noise_std < 0 

90 if np.any(select): 

91 LOGGER.warning("Negative noise std values detected, setting to 0") 

92 par.bkg_noise_std[select] = 0 

93 noise_draw *= par.bkg_noise_std 

94 noise_draw += par.bkg_noise_amp + variable_amp 

95 ctx.image += noise_draw 

96 ctx.image_mask = bkg_rms != 0 

97 

98 

99def add_from_chunked_map(ctx): 

100 """ 

101 Adds noise to the image using a rms map. The map is memmapped and only read out in 

102 quadratic chunks to reduce the memory footprint. 

103 

104 :param ctx.image: Simulated image 

105 :param ctx.params.maps_remote_dir: Path to fits-image containing the RMS of the 

106 noise in every pixel of the image 

107 :param ctx.params.bkg_rms_file_name: Path to fits-image containing the RMS of the 

108 noise in every pixel of the image 

109 :param ctx.params.bkg_noise_scale: Scale factor applied to the map 

110 :param ctx.params.bkg_noise_amp: Amplitude factor applied to the map 

111 """ 

112 par = ctx.parameters 

113 img_shape = ctx.image.shape 

114 

115 chunk = (par.chunksize, par.chunksize) # quadratic chunks for now 

116 

117 filename_h5, dataset = sysmaps_util.get_hdf_location_bgrms(par) 

118 filepath_h5 = io_util.get_abs_path(filename_h5, par.maps_remote_dir) 

119 

120 bkg_noise_factor = get_effective_bkg_noise_scale_factor(par) 

121 

122 with h5py.File(filepath_h5, "r") as file_bkgrms: 

123 bkg_rms = file_bkgrms[dataset] 

124 for i in range(img_shape[0] // chunk[0]): 

125 idx0 = slice(i * chunk[0], (i + 1) * chunk[0]) 

126 for j in range(img_shape[1] // chunk[1]): 

127 idx1 = slice(j * chunk[1], (j + 1) * chunk[1]) 

128 # TODO: implement variable systematics 

129 ctx.image[idx0, idx1] += ( 

130 par.bkg_noise_amp 

131 + bkg_noise_factor 

132 * par.bkg_noise_scale 

133 * np.random.normal(0, 1.0, size=chunk) 

134 * bkg_rms[idx0, idx1] 

135 ) 

136 

137 ctx.image_mask[idx0, idx1] = bkg_rms[idx0, idx1] != 0 

138 

139 

140BACKGROUND_GENERATOR = { 

141 "gaussian": add_from_gaussian, 

142 "map": add_from_map, 

143 "chunked_map": add_from_chunked_map, 

144} 

145 

146 

147class Plugin(BasePlugin): 

148 """ 

149 Generating random Gaussian background noise for each pixel and adding the noise onto 

150 the image. 

151 

152 :param background_sigma: rms spread of (zero mean) Gaussian background noise 

153 :param background_seed_offset: seed of the rng used to generate the background noies 

154 (optional) 

155 

156 :return: Image with noise added 

157 """ 

158 

159 def __call__(self): 

160 # Reseed random library 

161 np.random.seed( 

162 self.ctx.parameters.seed + self.ctx.parameters.background_seed_offset 

163 ) 

164 

165 # generate background noise 

166 generator = BACKGROUND_GENERATOR[self.ctx.parameters.background_type] 

167 generator(self.ctx) 

168 

169 def __str__(self): 

170 return "add background noise"