Coverage for src/ufig/plugins/render_galaxies_flexion.py: 96%

60 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-12 19:08 +0000

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""" 

8 

9import numba as nb 

10import numpy as np 

11from cosmic_toolbox.logger import get_logger 

12from ivy.plugin.base_plugin import BasePlugin 

13 

14from ufig import rendering_util 

15from ufig.plugins.gamma_interpolation_table import load_intrinsicTable 

16 

17LOGGER = get_logger(__file__) 

18 

19NAME = "render galaxies (flexion)" 

20intrinsicTableCache = None 

21 

22 

23@nb.jit(nopython=True, nogil=True) 

24def integrate_image( 

25 buffer, 

26 gammaprecision, 

27 gammaprecisionhigh, 

28 sersicprecision, 

29 intrinsic_table, 

30 image, 

31 x, 

32 y, 

33 nphot, 

34 sersic_indices, 

35 gal_ellip_matrices, 

36 psf_betas, 

37 psf_ellip_matrices, 

38 psf_flexion_tensors, 

39 psf_kurtoses, 

40 psf_dx_offset, 

41 psf_dy_offset, 

42 psf_flexion_suppression, 

43 sort_by_y, 

44): 

45 sin_table, cos_table = rendering_util.sin_cos_table() 

46 alpha_by_fwhm, beta_power = rendering_util.moffat_cdf_factors(psf_betas) 

47 

48 # sersic index lookup 

49 n = (sersic_indices * np.float64(np.int64(1) << 32) / 10.0).astype(np.uint32) 

50 sersic_table_l = (n >> np.uint32(32 - sersicprecision)).astype(np.uint32) 

51 sersic_table_b = (n - (sersic_table_l << np.uint32(32 - sersicprecision))).astype( 

52 np.float32 

53 ) 

54 sersic_table_b /= np.float32(1 << (32 - sersicprecision)) 

55 

56 ind_galaxies = np.argsort(y) if sort_by_y else np.arange(len(y)) 

57 

58 rng = 0 

59 

60 for i_gal in ind_galaxies: 

61 for i_beta in range(psf_betas.shape[1]): 

62 psf_flexion_suppress = rendering_util.psf_flexion_suppression( 

63 psf_betas[i_gal, i_beta], psf_flexion_suppression 

64 ) 

65 psf_kurtosis_suppress = rendering_util.psf_kurtosis_suppression( 

66 psf_betas[i_gal, i_beta] 

67 ) 

68 

69 for _ in range(nphot[i_gal, i_beta]): 

70 dr_psf, dx_psf, dy_psf = rendering_util.draw_photon_psf( 

71 buffer, 

72 sin_table, 

73 cos_table, 

74 alpha_by_fwhm[i_gal, i_beta], 

75 beta_power[i_gal, i_beta], 

76 rng, 

77 psf_ellip_matrices[i_gal], 

78 psf_flexion_tensors[i_gal], 

79 psf_kurtoses[i_gal, i_beta], 

80 psf_flexion_suppress, 

81 psf_kurtosis_suppress, 

82 psf_dx_offset[i_gal, i_beta], 

83 psf_dy_offset[i_gal, i_beta], 

84 ) 

85 

86 dx_gal, dy_gal = rendering_util.draw_photon_gal( 

87 buffer, 

88 sin_table, 

89 cos_table, 

90 gammaprecision, 

91 gammaprecisionhigh, 

92 sersic_table_l[i_gal], 

93 sersic_table_b[i_gal], 

94 intrinsic_table, 

95 rng, 

96 gal_ellip_matrices[i_gal], 

97 ) 

98 

99 dx = dx_psf + dx_gal 

100 dy = dy_psf + dy_gal 

101 

102 rendering_util.add_photon(image, x[i_gal], y[i_gal], dx, dy) 

103 

104 rng += 4 

105 

106 if rng + 4 > 44497: 

107 buffer[:21034] += buffer[(44497 - 21034) : 44497] 

108 buffer[21034:44497] += buffer[: (44497 - 21034)] 

109 rng = 0 

110 

111 

112class Plugin(BasePlugin): 

113 def __str__(self): 

114 return NAME 

115 

116 def __call__(self): 

117 par = self.ctx.parameters 

118 

119 LOGGER.info(f"Rendering {self.ctx.numgalaxies} galaxies") 

120 

121 # Seed 

122 np.random.seed(par.seed + par.gal_render_seed_offset) 

123 

124 # Intrinsic galaxy shapes 

125 gal_ellip_matrices = rendering_util.compute_ellip_matrices( 

126 self.ctx.galaxies.r50, self.ctx.galaxies.e1, self.ctx.galaxies.e2 

127 ) 

128 

129 # PSF properties 

130 if self.ctx.galaxies.psf_beta.shape[1] == 1: 

131 self.ctx.galaxies.psf_flux_ratio = np.ones(self.ctx.numgalaxies) 

132 

133 psf_ellip_matrices = rendering_util.compute_ellip_matrices( 

134 self.ctx.galaxies.psf_fwhm, 

135 self.ctx.galaxies.psf_e1, 

136 self.ctx.galaxies.psf_e2, 

137 ) 

138 

139 psf_flexion_tensors = rendering_util.compute_flexion_tensors( 

140 self.ctx.galaxies.psf_fwhm, 

141 self.ctx.galaxies.psf_f1, 

142 self.ctx.galaxies.psf_f2, 

143 self.ctx.galaxies.psf_g1, 

144 self.ctx.galaxies.psf_g2, 

145 ) 

146 

147 psf_kurtoses = rendering_util.compute_kurtoses( 

148 self.ctx.galaxies.psf_fwhm, 

149 self.ctx.galaxies.psf_kurtosis, 

150 self.ctx.galaxies.psf_beta, 

151 ) 

152 

153 # Calculate number of photons 

154 nphot = rendering_util.distribute_photons_psf_profiles( 

155 self.ctx.galaxies.nphot, 

156 self.ctx.galaxies.psf_beta.shape[1], 

157 self.ctx.galaxies.psf_flux_ratio, 

158 ) 

159 

160 # By default, if intrinsicTable is not loaded/computed, a corresponding one in 

161 # res/intrinsictables/ is loaded 

162 if not hasattr(self.ctx, "intrinsicTable"): 162 ↛ 167line 162 didn't jump to line 167 because the condition on line 162 was always true

163 sersic_table = load_intrinsicTable( 

164 par.sersicprecision, par.gammaprecision, par.gammaprecisionhigh 

165 ) 

166 else: 

167 sersic_table = self.ctx.intrinsicTable.copy() 

168 del self.ctx.intrinsicTable 

169 

170 if par.n_threads_photon_rendering > 1: 

171 ind_split = rendering_util.split_array( 

172 self.ctx.galaxies.x, par.n_threads_photon_rendering 

173 ) 

174 image = self.ctx.image 

175 self.ctx.image = None 

176 

177 rendering_args_split = [ 

178 [ 

179 rendering_util.rng_buffer() 

180 for _ in range(par.n_threads_photon_rendering) 

181 ], 

182 [par.gammaprecision * 1 for _ in range(par.n_threads_photon_rendering)], 

183 [ 

184 par.gammaprecisionhigh * 1 

185 for _ in range(par.n_threads_photon_rendering) 

186 ], 

187 [ 

188 par.sersicprecision * 1 

189 for _ in range(par.n_threads_photon_rendering) 

190 ], 

191 [sersic_table for _ in range(par.n_threads_photon_rendering)], 

192 [image for _ in range(par.n_threads_photon_rendering)], 

193 [self.ctx.galaxies.x[i_split] for i_split in ind_split], 

194 [self.ctx.galaxies.y[i_split] for i_split in ind_split], 

195 [nphot[i_split] for i_split in ind_split], 

196 [self.ctx.galaxies.sersic_n[i_split] for i_split in ind_split], 

197 [gal_ellip_matrices[i_split] for i_split in ind_split], 

198 [self.ctx.galaxies.psf_beta[i_split] for i_split in ind_split], 

199 [psf_ellip_matrices[i_split] for i_split in ind_split], 

200 [psf_flexion_tensors[i_split] for i_split in ind_split], 

201 [psf_kurtoses[i_split] for i_split in ind_split], 

202 [self.ctx.galaxies.psf_dx_offset[i_split] for i_split in ind_split], 

203 [self.ctx.galaxies.psf_dy_offset[i_split] for i_split in ind_split], 

204 [ 

205 par.psf_flexion_suppression 

206 for _ in range(par.n_threads_photon_rendering) 

207 ], 

208 [True for _ in range(par.n_threads_photon_rendering)], 

209 ] 

210 

211 # Run in threads 

212 rendering_util.execute_threaded( 

213 integrate_image, par.n_threads_photon_rendering, *rendering_args_split 

214 ) 

215 self.ctx.image = image 

216 

217 else: 

218 integrate_image( 

219 rendering_util.rng_buffer(), 

220 par.gammaprecision, 

221 par.gammaprecisionhigh, 

222 par.sersicprecision, 

223 sersic_table, 

224 self.ctx.image, 

225 self.ctx.galaxies.x, 

226 self.ctx.galaxies.y, 

227 nphot, 

228 self.ctx.galaxies.sersic_n, 

229 gal_ellip_matrices, 

230 self.ctx.galaxies.psf_beta, 

231 psf_ellip_matrices, 

232 psf_flexion_tensors, 

233 psf_kurtoses, 

234 self.ctx.galaxies.psf_dx_offset, 

235 self.ctx.galaxies.psf_dy_offset, 

236 par.psf_flexion_suppression, 

237 False, 

238 )