Coverage for src/ufig/rendering_util.py: 100%

152 statements  

« prev     ^ index     » next       coverage.py v7.10.2, created at 2025-08-07 15:17 +0000

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

2""" 

3Created on Jul 31, 2017 

4@author: Joerg Herbel 

5""" 

6 

7from concurrent.futures import ThreadPoolExecutor 

8 

9import numba as nb 

10import numpy as np 

11 

12# Gauss-Legendre nodes and weights 

13# Generated by http://keisan.casio.com/has10/SpecExec.cgi?id=system/2006/1280624821 

14GAUSS_LEGENDRE_NODES = np.array( 

15 [0.5 - 0.7745966692414833770359 / 2, 0.5, 0.5 + 0.7745966692414833770359 / 2], 

16 dtype=np.float32, 

17) 

18 

19GAUSS_LEGENDRE_WEIGHTS = np.array( 

20 [ 

21 0.555555555555555555556 / 2, 

22 0.8888888888888888888889 / 2, 

23 0.555555555555555555556 / 2, 

24 ], 

25 dtype=np.float32, 

26) 

27 

28GAUSS_LEGENDRE_WEIGHTS_2D = np.outer(GAUSS_LEGENDRE_WEIGHTS, GAUSS_LEGENDRE_WEIGHTS) 

29 

30 

31def rng_buffer(): 

32 buffer = np.random.randint(0, 1 << 32, size=(44497,)).astype(np.uint32) 

33 return buffer 

34 

35 

36@nb.jit(nopython=True) 

37def sin_cos_table(): 

38 r = ( 

39 np.arange(0, (1 << 11) + 1).astype(np.float64) 

40 * 2.0 

41 * np.pi 

42 / np.float64(1 << 11) 

43 ) 

44 sin_table = np.sin(r) 

45 cos_table = np.cos(r) 

46 return sin_table, cos_table 

47 

48 

49def compute_ellip_matrices(size, e1, e2): 

50 """ 

51 size can either be the PSF fwhm or the intrinsic galaxy r50 

52 """ 

53 

54 m = np.empty((len(size), 2, 2)) 

55 m[:, 0, 0] = 1 

56 m[:, 0, 1] = 0 

57 m[:, 1, 0] = 0 

58 m[:, 1, 1] = 1 

59 

60 e = np.sqrt(e1**2 + e2**2) 

61 mask_e = e > 1e-7 

62 

63 m[mask_e, 0, 0] = np.sqrt( 

64 (1.0 + e[mask_e]) * (1.0 + e1[mask_e] / e[mask_e]) / 2.0 

65 ) * np.copysign(1.0, e2[mask_e]) 

66 m[mask_e, 0, 1] = np.sqrt((1.0 - e[mask_e]) * (1.0 - e1[mask_e] / e[mask_e]) / 2.0) 

67 m[mask_e, 1, 0] = np.sqrt((1.0 + e[mask_e]) * (1.0 - e1[mask_e] / e[mask_e]) / 2.0) 

68 m[mask_e, 1, 1] = np.sqrt( 

69 (1.0 - e[mask_e]) * (1.0 + e1[mask_e] / e[mask_e]) / 2.0 

70 ) * np.copysign(1.0, e2[mask_e]) 

71 

72 m *= size[:, None, None] 

73 

74 return m 

75 

76 

77def compute_flexion_tensors(fwhm, f1, f2, g1, g2): 

78 d = np.zeros((len(fwhm), 2, 2, 2)) 

79 

80 d[:, 0, 0, 0] = 3 * f1 + g1 

81 d[:, 0, 1, 0] = 2 * (f2 + g2) 

82 # we deliberately do not compute this entry, since it equals d[:, 0, 1, 0]; 

83 # instead, we use d[:, 0, 1, 0] = d[:, 0, 1, 0] + d[:, 0, 0, 1] 

84 # d[:, 0, 0, 1] = f2 + g2 

85 d[:, 0, 1, 1] = f1 - g1 

86 

87 d[:, 1, 0, 0] = f2 + g2 

88 d[:, 1, 1, 0] = 2 * (f1 - g1) 

89 # see comment for d[:, 0, 0, 1] 

90 # d[:, 1, 0, 1] = f1 - g1 

91 d[:, 1, 1, 1] = 3 * f2 - g2 

92 

93 d *= -0.5 * fwhm[:, None, None, None] 

94 

95 return d 

96 

97 

98def compute_kurtoses(fwhm, kurtoses, beta): 

99 k = kurtoses[:, np.newaxis] / beta * fwhm[:, np.newaxis] 

100 return k 

101 

102 

103@nb.jit(nopython=True) 

104def psf_flexion_suppression(beta, flexion_suppression): 

105 suppress_factor = flexion_suppression / beta**2.0 

106 return suppress_factor 

107 

108 

109@nb.jit(nopython=True) 

110def psf_kurtosis_suppression(beta): 

111 suppress_factor = -12.5 / beta**2.0 

112 return suppress_factor 

113 

114 

115def distribute_photons_psf_profiles(nphot, n_profiles, flux_ratio): 

116 nphot_split = np.zeros((len(nphot), n_profiles), dtype=np.int64) 

117 flux_ratio_clip = np.clip(flux_ratio, a_min=0, a_max=1) 

118 nphot_split[:, 0] = np.array(np.round(flux_ratio_clip * nphot), dtype=np.int64) 

119 

120 for i in range(1, n_profiles): 

121 nphot_split[:, i] = nphot - nphot_split[:, i - 1] 

122 

123 return nphot_split 

124 

125 

126@nb.jit(nopython=True) 

127def moffat_cdf_factors(psf_beta): 

128 alpha_by_fwhm = 1.0 / (2.0 * np.sqrt(2.0 ** (1.0 / psf_beta) - 1.0)) 

129 beta_power = 1.0 / (1.0 - psf_beta) 

130 return alpha_by_fwhm, beta_power 

131 

132 

133@nb.jit(nopython=True) 

134def draw_photon_psf_radial( 

135 rng_buffer, sin_table, cos_table, rng, alpha_by_fwhm, beta_power 

136): 

137 # sin/cos 

138 scL = rng_buffer[rng] >> (32 - 11) 

139 scB = np.float64(rng_buffer[rng] & np.uint32((1 << (32 - 11)) - 1)) / np.float64( 

140 1 << (32 - 11) 

141 ) 

142 scA = 1.0 - scB 

143 sn = scA * sin_table[scL] + scB * sin_table[scL + 1] 

144 cn = scA * cos_table[scL] + scB * cos_table[scL + 1] 

145 

146 # dr 

147 cdf = np.float64( 

148 rng_buffer[rng + 1] & np.uint32((1 << (32 - 11)) - 1) 

149 ) / np.float64(1 << (32 - 11)) 

150 dr = alpha_by_fwhm * np.sqrt((1 - cdf) ** beta_power - 1) 

151 

152 return dr, sn, cn 

153 

154 

155@nb.jit(nopython=True) 

156def draw_photon_psf( 

157 rng_buffer, 

158 sin_table, 

159 cos_table, 

160 alpha_by_fwhm, 

161 beta_power, 

162 rng, 

163 ellip_matrix, 

164 flexion_tensor, 

165 kurtosis, 

166 flexion_suppression, 

167 kurtosis_suppression, 

168 dx_offset, 

169 dy_offset, 

170): 

171 dr, sn, cn = draw_photon_psf_radial( 

172 rng_buffer, sin_table, cos_table, rng, alpha_by_fwhm, beta_power 

173 ) 

174 

175 # update position 

176 x_phot_circ = dr * cn 

177 y_phot_circ = dr * sn 

178 

179 # ellipticity 

180 dx = ellip_matrix[0, 0] * x_phot_circ - ellip_matrix[0, 1] * y_phot_circ 

181 dy = ellip_matrix[1, 0] * x_phot_circ + ellip_matrix[1, 1] * y_phot_circ 

182 

183 # flexion 

184 x_phot_circ_sq = x_phot_circ**2 

185 y_phot_circ_sq = y_phot_circ**2 

186 r_phot_sq = x_phot_circ_sq + y_phot_circ_sq 

187 x_phot_y_phot_circ = x_phot_circ * y_phot_circ 

188 

189 flexion_suppress = np.exp(flexion_suppression * r_phot_sq) 

190 

191 dx += ( 

192 flexion_tensor[0, 0, 0] * x_phot_circ_sq 

193 + flexion_tensor[0, 1, 0] * x_phot_y_phot_circ 

194 + flexion_tensor[0, 1, 1] * y_phot_circ_sq 

195 ) * flexion_suppress 

196 

197 dy += ( 

198 flexion_tensor[1, 0, 0] * x_phot_circ_sq 

199 + flexion_tensor[1, 1, 0] * x_phot_y_phot_circ 

200 + flexion_tensor[1, 1, 1] * y_phot_circ_sq 

201 ) * flexion_suppress 

202 

203 # kurtosis 

204 kurtosis_suppress = np.exp(kurtosis_suppression * r_phot_sq) 

205 dx += (kurtosis * r_phot_sq * x_phot_circ) * kurtosis_suppress 

206 dy += (kurtosis * r_phot_sq * y_phot_circ) * kurtosis_suppress 

207 

208 # Offset between multiple profiles 

209 dx += dx_offset 

210 dy += dy_offset 

211 

212 return dr, dx, dy 

213 

214 

215@nb.jit(nopython=True) 

216def draw_photon_gal( 

217 rng_buffer, 

218 sin_table, 

219 cos_table, 

220 gammaprecision, 

221 gammaprecisionhigh, 

222 intrinsic_sersic_l, 

223 intrinsic_sersic_b, 

224 intrinsic_table, 

225 rng, 

226 ellip_matrix, 

227): 

228 # sin/cos 

229 scL = rng_buffer[rng + 2] >> (32 - 11) 

230 scB = np.float64( 

231 rng_buffer[rng + 2] & np.uint32((1 << (32 - 11)) - 1) 

232 ) / np.float64(1 << (32 - 11)) 

233 scA = 1.0 - scB 

234 sn = scA * sin_table[scL] + scB * sin_table[scL + 1] 

235 cn = scA * cos_table[scL] + scB * cos_table[scL + 1] 

236 

237 # dr 

238 drMask = ( 

239 rng_buffer[rng + 3] >> (32 - gammaprecisionhigh) 

240 == (1 << gammaprecisionhigh) - 1 

241 ) 

242 drK = rng_buffer[rng + 3] << np.uint32(drMask * gammaprecisionhigh) 

243 drL = ((drK >> (32 - gammaprecision)) & ((1 << gammaprecision) - 1)) + np.uint32( 

244 drMask * (1 << gammaprecision) 

245 ) 

246 drB = np.float64( 

247 rng_buffer[rng + 3] 

248 & ((1 << (32 - gammaprecision - np.uint32(drMask * gammaprecisionhigh))) - 1) 

249 ) / np.float64(1 << (32 - gammaprecision - np.uint32(drMask * gammaprecisionhigh))) 

250 drA = 1.0 - drB 

251 

252 nL = intrinsic_sersic_l 

253 nB = intrinsic_sersic_b 

254 nA = 1 - nB 

255 

256 dr = drA * ( 

257 nA * intrinsic_table[nL, drL] + nB * intrinsic_table[nL + 1, drL] 

258 ) + drB * ( 

259 nA * intrinsic_table[nL, drL + 1] + nB * intrinsic_table[nL + 1, drL + 1] 

260 ) 

261 

262 dx = dr * (cn * ellip_matrix[0, 0] - sn * ellip_matrix[0, 1]) 

263 dy = dr * (cn * ellip_matrix[1, 0] + sn * ellip_matrix[1, 1]) 

264 

265 return dx, dy 

266 

267 

268@nb.jit(nopython=True) 

269def add_photon(image, x, y, dx, dy): 

270 """ 

271 Add one photon to image. 

272 

273 :param image: image 

274 :param x: centroid x-position 

275 :param y: centroid y-position 

276 :param dx: x-coordinate sampled from light profile 

277 :param dy: y-coordinate sampled from light profile 

278 :return: 

279 """ 

280 x_dx = x + dx 

281 y_dy = y + dy 

282 if 0 <= x_dx < image.shape[1] and 0 <= y_dy < image.shape[0]: 

283 image[np.int32(y_dy), np.int32(x_dx)] += 1 

284 

285 

286@nb.jit(nopython=True) 

287def integrate_pixel_psf( 

288 alpha_sq, 

289 beta, 

290 r50, 

291 e1, 

292 e2, 

293 dx_offset, 

294 dy_offset, 

295 nphot, 

296 min_x, 

297 max_x, 

298 min_y, 

299 max_y, 

300 offset_x, 

301 offset_y, 

302 mean, 

303): 

304 e1_plus = 1 + e1 

305 e1_minus = 1 - e1 

306 e2_2 = 2 * e2 

307 normalization_sq = 1 - (e1**2 + e2**2) 

308 

309 r50_sq = r50**2 

310 r50_sq_alpha_sq = r50_sq * alpha_sq 

311 fac = (beta - 1) / np.sqrt(normalization_sq) / (np.pi * r50_sq_alpha_sq) 

312 

313 for x in range(min_x, max_x + 1): 

314 xdiff = x - offset_x 

315 

316 for y in range(min_y, max_y + 1): 

317 density = 0.0 

318 ydiff = y - offset_y 

319 

320 for beta_ind in range(len(beta)): 

321 for xi in range(0, len(GAUSS_LEGENDRE_WEIGHTS)): 

322 x_centered = xdiff + GAUSS_LEGENDRE_NODES[xi] - dx_offset[beta_ind] 

323 x_centered_e1 = x_centered**2 * e1_minus 

324 x_centered_e2_2 = e2_2 * x_centered 

325 

326 for yi in range(0, len(GAUSS_LEGENDRE_WEIGHTS)): 

327 y_centered = ( 

328 ydiff + GAUSS_LEGENDRE_NODES[yi] - dy_offset[beta_ind] 

329 ) 

330 dr_sq = ( 

331 x_centered_e1 

332 + y_centered**2 * e1_plus 

333 - x_centered_e2_2 * y_centered 

334 ) / normalization_sq 

335 density += ( 

336 nphot[beta_ind] 

337 * GAUSS_LEGENDRE_WEIGHTS_2D[yi, xi] 

338 * fac[beta_ind] 

339 * (1 + dr_sq / r50_sq_alpha_sq[beta_ind]) 

340 ** (-beta[beta_ind]) 

341 ) 

342 

343 mean[y - min_y, x - min_x] = np.random.poisson(density) 

344 

345 return mean 

346 

347 

348def split_array(arr, n_split): 

349 x_sort_ind = np.argsort(arr) 

350 ind_split = np.linspace(0, len(arr), num=n_split + 1, dtype=np.int32) 

351 

352 indices_split = [None] * n_split 

353 for i_split in range(n_split): 

354 indices_split[i_split] = x_sort_ind[ind_split[i_split] : ind_split[i_split + 1]] 

355 

356 return indices_split 

357 

358 

359def execute_threaded(func, n_threads, *args): 

360 with ThreadPoolExecutor(max_workers=n_threads) as executor: 

361 executor.map(func, *args)