Coverage for src/ufig/plugins/render_stars_photon.py: 92%

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

2""" 

3Created on Jul 31, 2017 

4@author: Joerg Herbel 

5""" 

6 

7import numba as nb 

8import numpy as np 

9from cosmic_toolbox.logger import get_logger 

10from ivy.plugin.base_plugin import BasePlugin 

11 

12from ufig import rendering_util 

13from ufig.plugins import add_psf 

14 

15LOGGER = get_logger(__file__) 

16 

17NAME = "render stars photon-based" 

18 

19 

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

21def integrate_image_phot( 

22 buffer, 

23 image, 

24 x, 

25 y, 

26 nphot, 

27 psf_betas, 

28 ellip_matrices, 

29 flexion_tensors, 

30 kurtoses, 

31 dx_offset, 

32 dy_offset, 

33 flexion_suppression, 

34 sort_by_y, 

35): 

36 sin_table, cos_table = rendering_util.sin_cos_table() 

37 alpha_by_fwhm, beta_power = rendering_util.moffat_cdf_factors(psf_betas) 

38 

39 ind_stars = np.argsort(y) if sort_by_y else np.arange(len(y)) 

40 

41 rng = 0 

42 

43 for i_star in ind_stars: 

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

45 flexion_suppress = rendering_util.psf_flexion_suppression( 

46 psf_betas[i_star, i_beta], flexion_suppression 

47 ) 

48 kurtosis_suppress = rendering_util.psf_kurtosis_suppression( 

49 psf_betas[i_star, i_beta] 

50 ) 

51 

52 for _ in range(nphot[i_star, i_beta]): 

53 dr, dx, dy = rendering_util.draw_photon_psf( 

54 buffer, 

55 sin_table, 

56 cos_table, 

57 alpha_by_fwhm[i_star, i_beta], 

58 beta_power[i_star, i_beta], 

59 rng, 

60 ellip_matrices[i_star], 

61 flexion_tensors[i_star], 

62 kurtoses[i_star, i_beta], 

63 flexion_suppress, 

64 kurtosis_suppress, 

65 dx_offset[i_star, i_beta], 

66 dy_offset[i_star, i_beta], 

67 ) 

68 

69 rendering_util.add_photon(image, x[i_star], y[i_star], dx, dy) 

70 

71 rng += 2 

72 if rng + 2 > 44497: 

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

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

75 rng = 0 

76 

77 

78@nb.jit(nopython=True) 

79def pixel_integration( 

80 seed, 

81 image, 

82 r_max, 

83 x, 

84 y, 

85 nphot, 

86 psf_beta, 

87 psf_r50, 

88 psf_e1, 

89 psf_e2, 

90 psf_dx_offset, 

91 psf_dy_offset, 

92): 

93 # Seed within numba-compiled code, otherwise this code won't be affected 

94 np.random.seed(seed) 

95 

96 size_y, size_x = image.shape 

97 alpha_sq = 1.0 / (2.0 ** (1.0 / (psf_beta - 1.0)) - 1.0) 

98 

99 # bounds for x axis 

100 center_x = x.astype(np.int_) 

101 offset_x = x - np.floor(x) 

102 min_x = -r_max 

103 idx = center_x < r_max 

104 min_x[idx] = -center_x[idx] 

105 max_x = ( 

106 r_max.copy() 

107 ) # otherwise it will be a pointer, which overwritten and then re-used 

108 idx = size_x - r_max <= center_x 

109 max_x[idx] = size_x - 1 - center_x[idx] 

110 

111 # bounds for y axis 

112 center_y = y.astype(np.int_) 

113 offset_y = y - np.floor(y) 

114 min_y = -r_max 

115 idy = center_y < r_max 

116 min_y[idy] = -center_y[idy] 

117 max_y = ( 

118 r_max.copy() 

119 ) # otherwise it will be a pointer, which overwritten and then re-used 

120 idy = size_y - r_max <= center_y 

121 max_y[idy] = size_y - 1 - center_y[idy] 

122 

123 for i in range(len(x)): 

124 star = np.zeros( 

125 (max_y[i] - min_y[i] + 1, max_x[i] - min_x[i] + 1), dtype=image.dtype 

126 ) 

127 

128 if star.shape == (1, 1): 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true

129 continue 

130 

131 rendering_util.integrate_pixel_psf( 

132 alpha_sq[i], 

133 psf_beta[i], 

134 psf_r50[i], 

135 psf_e1[i], 

136 psf_e2[i], 

137 psf_dx_offset[i], 

138 psf_dy_offset[i], 

139 nphot[i], 

140 min_x[i], 

141 max_x[i], 

142 min_y[i], 

143 max_y[i], 

144 offset_x[i], 

145 offset_y[i], 

146 star, 

147 ) 

148 

149 image[ 

150 min_y[i] + center_y[i] : max_y[i] + 1 + center_y[i], 

151 min_x[i] + center_x[i] : max_x[i] + 1 + center_x[i], 

152 ] += star 

153 

154 

155def integrate_image_pixel( 

156 seed, 

157 image, 

158 gain, 

159 render_stars_accuracy, 

160 x, 

161 y, 

162 nphot, 

163 psf_beta, 

164 psf_flux_ratio, 

165 psf_fwhm, 

166 psf_e1, 

167 psf_e2, 

168 psf_dx_offset, 

169 psf_dy_offset, 

170): 

171 # integrate a bigger radius for brighter stars 

172 index_min = np.argmin(psf_beta, axis=1) 

173 beta_min = psf_beta[np.arange(psf_beta.shape[0]), index_min] 

174 

175 select = index_min > 0 

176 flux_ratio_min = psf_flux_ratio.copy() 

177 flux_ratio_min[select] = 1 - flux_ratio_min[select] 

178 

179 alpha_min = add_psf.moffat_fwhm2alpha(psf_fwhm, beta_min) 

180 flux = ( 

181 nphot 

182 * flux_ratio_min 

183 * (beta_min - 1) 

184 / (np.pi * gain) 

185 * alpha_min ** (2 * (beta_min - 1)) 

186 ) 

187 r_max = ((flux / render_stars_accuracy) ** (1.0 / (2 * beta_min))).astype(np.int_) 

188 r_max += 2 

189 

190 # Distribute photons to individual Moffat profiles 

191 nphot_split = rendering_util.distribute_photons_psf_profiles( 

192 nphot, psf_beta.shape[1], psf_flux_ratio 

193 ) 

194 

195 # Get r50 

196 psf_r50 = np.empty_like(psf_beta) 

197 

198 for beta_ind in range(psf_beta.shape[1]): 

199 psf_r50[:, beta_ind] = add_psf.moffat_fwhm2r50( 

200 psf_fwhm, psf_beta[:, beta_ind], psf_flux_ratio 

201 ) 

202 

203 pixel_integration( 

204 seed, 

205 image, 

206 r_max, 

207 x, 

208 y, 

209 nphot_split, 

210 psf_beta, 

211 psf_r50, 

212 psf_e1, 

213 psf_e2, 

214 psf_dx_offset, 

215 psf_dy_offset, 

216 ) 

217 

218 

219@nb.jit(nopython=True) 

220def integrate_cube( 

221 buffer, 

222 cube, 

223 x, 

224 y, 

225 nphot, 

226 psf_betas, 

227 ellip_matrices, 

228 flexion_tensors, 

229 kurtoses, 

230 dx_offset, 

231 dy_offset, 

232 flexion_suppression, 

233 q_xx, 

234 q_yy, 

235 q_xy, 

236): 

237 sin_table, cos_table = rendering_util.sin_cos_table() 

238 alpha_by_fwhm, beta_power = rendering_util.moffat_cdf_factors(psf_betas) 

239 

240 rng = 0 

241 cov_cutoff = 5 

242 

243 for i_star in range(x.size): 

244 n_phot_cov = 0 

245 mean_phot_x = 0.0 

246 mean_phot_y = 0.0 

247 

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

249 flexion_suppress = rendering_util.psf_flexion_suppression( 

250 psf_betas[i_star, i_beta], flexion_suppression 

251 ) 

252 kurtosis_suppress = rendering_util.psf_kurtosis_suppression( 

253 psf_betas[i_star, i_beta] 

254 ) 

255 

256 for _ in range(nphot[i_star, i_beta]): 

257 dr, dx, dy = rendering_util.draw_photon_psf( 

258 buffer, 

259 sin_table, 

260 cos_table, 

261 alpha_by_fwhm[i_star, i_beta], 

262 beta_power[i_star, i_beta], 

263 rng, 

264 ellip_matrices[i_star], 

265 flexion_tensors[i_star], 

266 kurtoses[i_star, i_beta], 

267 flexion_suppress, 

268 kurtosis_suppress, 

269 dx_offset[i_star, i_beta], 

270 dy_offset[i_star, i_beta], 

271 ) 

272 

273 rendering_util.add_photon(cube[i_star], x[i_star], y[i_star], dx, dy) 

274 

275 rng += 2 

276 if rng + 2 > 44497: 276 ↛ 277line 276 didn't jump to line 277 because the condition on line 276 was never true

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

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

279 rng = 0 

280 

281 # moments 

282 if dr < cov_cutoff: 

283 n_phot_cov += 1 

284 delta_x = dx - mean_phot_x 

285 delta_y = dy - mean_phot_y 

286 mean_phot_x += delta_x / n_phot_cov 

287 mean_phot_y += delta_y / n_phot_cov 

288 delta_y_2 = dy - mean_phot_y 

289 q_xx[i_star] += delta_x * (dx - mean_phot_x) 

290 q_yy[i_star] += delta_y * delta_y_2 

291 q_xy[i_star] += delta_x * delta_y_2 

292 

293 if n_phot_cov > 1: 293 ↛ 243line 293 didn't jump to line 243 because the condition on line 293 was always true

294 q_xx[i_star] /= n_phot_cov - 1 

295 q_yy[i_star] /= n_phot_cov - 1 

296 q_xy[i_star] /= n_phot_cov - 1 

297 

298 

299def integrate_threaded( 

300 ctx, 

301 x, 

302 y, 

303 nphot, 

304 psf_beta, 

305 psf_ellip_matrices, 

306 psf_flexion_tensors, 

307 psf_kurtoses, 

308 psf_dx_offset, 

309 psf_dy_offset, 

310 flexion_suppression, 

311): 

312 n_threads = ctx.parameters.n_threads_photon_rendering 

313 

314 # Split workload between threads 

315 ind_split = rendering_util.split_array(x, n_threads) 

316 image = ctx.image 

317 ctx.image = None 

318 

319 rendering_args_split = [ 

320 [rendering_util.rng_buffer() for _ in range(n_threads)], 

321 [image for _ in range(n_threads)], 

322 [x[i_split] for i_split in ind_split], 

323 [y[i_split] for i_split in ind_split], 

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

325 [psf_beta[i_split] for i_split in ind_split], 

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

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

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

329 [psf_dx_offset[i_split] for i_split in ind_split], 

330 [psf_dy_offset[i_split] for i_split in ind_split], 

331 [flexion_suppression for _ in range(n_threads)], 

332 [True for _ in range(n_threads)], 

333 ] 

334 

335 # Run in threads 

336 # print('Number of threads used for rendering stars: {}'.format(n_threads)) 

337 rendering_util.execute_threaded( 

338 integrate_image_phot, n_threads, *rendering_args_split 

339 ) 

340 ctx.image = image 

341 

342 

343class Plugin(BasePlugin): 

344 """ 

345 Render stellar profiles photon-by-photon onto a pixelated grid. 

346 """ 

347 

348 def __str__(self): 

349 return NAME 

350 

351 def __call__(self): 

352 par = self.ctx.parameters 

353 

354 LOGGER.info(f"Rendering {self.ctx.numstars} stars") 

355 

356 # Seed 

357 current_seed = par.seed + par.star_render_seed_offset 

358 np.random.seed(par.seed + par.star_render_seed_offset) 

359 

360 # PSF properties 

361 if self.ctx.stars.psf_beta.shape[1] == 1: 

362 self.ctx.stars.psf_flux_ratio = np.ones(self.ctx.numstars) 

363 

364 psf_ellip_matrices = rendering_util.compute_ellip_matrices( 

365 self.ctx.stars.psf_fwhm, self.ctx.stars.psf_e1, self.ctx.stars.psf_e2 

366 ) 

367 

368 psf_flexion_tensors = rendering_util.compute_flexion_tensors( 

369 self.ctx.stars.psf_fwhm, 

370 self.ctx.stars.psf_f1, 

371 self.ctx.stars.psf_f2, 

372 self.ctx.stars.psf_g1, 

373 self.ctx.stars.psf_g2, 

374 ) 

375 

376 psf_kurtoses = rendering_util.compute_kurtoses( 

377 self.ctx.stars.psf_fwhm, 

378 self.ctx.stars.psf_kurtosis, 

379 self.ctx.stars.psf_beta, 

380 ) 

381 

382 # Calculate number of photons 

383 nphot = rendering_util.distribute_photons_psf_profiles( 

384 self.ctx.stars.nphot, 

385 self.ctx.stars.psf_beta.shape[1], 

386 self.ctx.stars.psf_flux_ratio, 

387 ) 

388 

389 if hasattr(self.ctx, "image"): 

390 if hasattr(par, "mag_pixel_rendering_stars"): 

391 mask_pix = self.ctx.stars.mag < par.mag_pixel_rendering_stars 

392 mask_phot = ~mask_pix 

393 

394 if np.any(mask_phot): 394 ↛ 395line 394 didn't jump to line 395 because the condition on line 394 was never true

395 if par.n_threads_photon_rendering > 1: 

396 integrate_threaded( 

397 self.ctx, 

398 self.ctx.stars.x[mask_phot], 

399 self.ctx.stars.y[mask_phot], 

400 nphot[mask_phot], 

401 self.ctx.stars.psf_beta[mask_phot], 

402 psf_ellip_matrices[mask_phot], 

403 psf_flexion_tensors[mask_phot], 

404 psf_kurtoses[mask_phot], 

405 self.ctx.stars.psf_dx_offset[mask_phot], 

406 self.ctx.stars.psf_dy_offset[mask_phot], 

407 par.psf_flexion_suppression, 

408 ) 

409 

410 else: 

411 integrate_image_phot( 

412 rendering_util.rng_buffer(), 

413 self.ctx.image, 

414 self.ctx.stars.x[mask_phot], 

415 self.ctx.stars.y[mask_phot], 

416 nphot[mask_phot], 

417 self.ctx.stars.psf_beta[mask_phot], 

418 psf_ellip_matrices[mask_phot], 

419 psf_flexion_tensors[mask_phot], 

420 psf_kurtoses[mask_phot], 

421 self.ctx.stars.psf_dx_offset[mask_phot], 

422 self.ctx.stars.psf_dy_offset[mask_phot], 

423 par.psf_flexion_suppression, 

424 False, 

425 ) 

426 

427 if np.any(mask_pix): 

428 integrate_image_pixel( 

429 current_seed, 

430 self.ctx.image, 

431 par.gain, 

432 par.render_stars_accuracy, 

433 self.ctx.stars.x[mask_pix], 

434 self.ctx.stars.y[mask_pix], 

435 self.ctx.stars.nphot[mask_pix], 

436 self.ctx.stars.psf_beta[mask_pix], 

437 self.ctx.stars.psf_flux_ratio[mask_pix], 

438 self.ctx.stars.psf_fwhm[mask_pix], 

439 self.ctx.stars.psf_e1[mask_pix], 

440 self.ctx.stars.psf_e2[mask_pix], 

441 self.ctx.stars.psf_dx_offset[mask_pix], 

442 self.ctx.stars.psf_dy_offset[mask_pix], 

443 ) 

444 

445 else: 

446 if par.n_threads_photon_rendering > 1: 

447 integrate_threaded( 

448 self.ctx, 

449 self.ctx.stars.x, 

450 self.ctx.stars.y, 

451 nphot, 

452 self.ctx.stars.psf_beta, 

453 psf_ellip_matrices, 

454 psf_flexion_tensors, 

455 psf_kurtoses, 

456 self.ctx.stars.psf_dx_offset, 

457 self.ctx.stars.psf_dy_offset, 

458 par.psf_flexion_suppression, 

459 ) 

460 

461 else: 

462 integrate_image_phot( 

463 rendering_util.rng_buffer(), 

464 self.ctx.image, 

465 self.ctx.stars.x, 

466 self.ctx.stars.y, 

467 nphot, 

468 self.ctx.stars.psf_beta, 

469 psf_ellip_matrices, 

470 psf_flexion_tensors, 

471 psf_kurtoses, 

472 self.ctx.stars.psf_dx_offset, 

473 self.ctx.stars.psf_dy_offset, 

474 par.psf_flexion_suppression, 

475 False, 

476 ) 

477 

478 elif hasattr(self.ctx, "star_cube"): 478 ↛ 503line 478 didn't jump to line 503 because the condition on line 478 was always true

479 # Initialize moments 

480 self.ctx.stars.q_xx = np.zeros(self.ctx.numstars) 

481 self.ctx.stars.q_yy = np.zeros(self.ctx.numstars) 

482 self.ctx.stars.q_xy = np.zeros(self.ctx.numstars) 

483 

484 integrate_cube( 

485 rendering_util.rng_buffer(), 

486 self.ctx.star_cube, 

487 self.ctx.stars.x, 

488 self.ctx.stars.y, 

489 nphot, 

490 self.ctx.stars.psf_beta, 

491 psf_ellip_matrices, 

492 psf_flexion_tensors, 

493 psf_kurtoses, 

494 self.ctx.stars.psf_dx_offset, 

495 self.ctx.stars.psf_dy_offset, 

496 par.psf_flexion_suppression, 

497 self.ctx.stars.q_xx, 

498 self.ctx.stars.q_yy, 

499 self.ctx.stars.q_xy, 

500 ) 

501 

502 else: 

503 raise RuntimeError( 

504 "No known photon container (image or star_cube) provided in ivy context" 

505 )