Coverage for yasfpy/optics.py: 0%

1import logging 

2from typing import Union 

3from math import ceil 

4 

5import numpy as np 

6from numba import cuda 

7from yasfpy.functions.spherical_functions_trigon import spherical_functions_trigon 

8 

9from yasfpy.simulation import Simulation 

10 

11from yasfpy.functions.cpu_numba import ( 

12 compute_scattering_cross_section, 

13 compute_electric_field_angle_components, 

14 compute_polarization_components, 

15) 

16from yasfpy.functions.cuda_numba import ( 

17 compute_scattering_cross_section_gpu, 

18 compute_electric_field_angle_components_gpu, 

19 compute_polarization_components_gpu, 

20) 

21 

22# from yasfpy.functions.cpu_numba import compute_scattering_cross_section, compute_radial_independent_scattered_field, compute_electric_field_angle_components, compute_polarization_components 

23# from yasfpy.functions.cuda_numba import compute_scattering_cross_section_gpu, compute_radial_independent_scattered_field_gpu, compute_electric_field_angle_components_gpu, compute_polarization_components_gpu 

24 

25 

26class Optics: 

27 """The above code is a Python program that is likely related to optics. However, without seeing 

28 the actual code, it is difficult to determine exactly what it is doing. 

29 """ 

30 

31 def __init__(self, simulation: Simulation): 

32 """ 

33 Initializes the Optics object. 

34 

35 Args: 

36 simulation (Simulation): An instance of the Simulation class. 

37 

38 """ 

39 self.simulation = simulation 

40 

41 self.c_ext = np.zeros_like(simulation.parameters.wavelength, dtype=complex) 

42 self.c_sca = np.zeros_like(simulation.parameters.wavelength, dtype=complex) 

43 

44 self.log = logging.getLogger(__name__) 

45 

46 def compute_cross_sections(self): 

47 """ 

48 Compute the scattering and extinction cross sections. 

49 

50 This method calculates the scattering and extinction cross sections for the simulation. 

51 It uses the initial field coefficients and scattered field coefficients to perform the calculations. 

52 """ 

53 a = self.simulation.initial_field_coefficients 

54 f = self.simulation.scattered_field_coefficients 

55 

56 self.c_ext = np.zeros( 

57 self.simulation.parameters.wavelengths_number, dtype=complex 

58 ) 

59 self.c_ext -= ( 

60 np.sum(np.conjugate(a) * f, axis=(0, 1)) 

61 / np.power(self.simulation.parameters.k_medium, 2) 

62 * np.pi 

63 ) 

64 

65 lmax = self.simulation.numerics.lmax 

66 particle_number = self.simulation.parameters.particles.number 

67 wavelengths = self.simulation.parameters.k_medium.shape[0] 

68 translation_table = self.simulation.numerics.translation_ab5 

69 associated_legendre_lookup = self.simulation.plm 

70 spherical_bessel_lookup = self.simulation.sph_j 

71 e_j_dm_phi_loopup = self.simulation.e_j_dm_phi 

72 

73 idx_lookup = self.simulation.idx_lookup 

74 

75 if self.simulation.numerics.gpu: 

76 c_sca_real = np.zeros(wavelengths, dtype=float) 

77 c_sca_imag = np.zeros_like(c_sca_real) 

78 

79 idx_device = cuda.to_device(idx_lookup) 

80 sfc_device = cuda.to_device(f) 

81 c_sca_real_device = cuda.to_device(c_sca_real) 

82 c_sca_imag_device = cuda.to_device(c_sca_imag) 

83 translation_device = cuda.to_device(translation_table) 

84 associated_legendre_device = cuda.to_device(associated_legendre_lookup) 

85 spherical_bessel_device = cuda.to_device(spherical_bessel_lookup) 

86 e_j_dm_phi_device = cuda.to_device(e_j_dm_phi_loopup) 

87 

88 jmax = particle_number * 2 * lmax * (lmax + 2) 

89 threads_per_block = (16, 16, 2) 

90 blocks_per_grid_x = ceil(jmax / threads_per_block[0]) 

91 blocks_per_grid_y = ceil(jmax / threads_per_block[1]) 

92 blocks_per_grid_z = ceil(wavelengths / threads_per_block[2]) 

93 blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y, blocks_per_grid_z) 

94 

95 compute_scattering_cross_section_gpu[blocks_per_grid, threads_per_block]( 

96 lmax, 

97 particle_number, 

98 idx_device, 

99 sfc_device, 

100 translation_device, 

101 associated_legendre_device, 

102 spherical_bessel_device, 

103 e_j_dm_phi_device, 

104 c_sca_real_device, 

105 c_sca_imag_device, 

106 ) 

107 c_sca_real = c_sca_real_device.copy_to_host() 

108 c_sca_imag = c_sca_imag_device.copy_to_host() 

109 c_sca = c_sca_real + 1j * c_sca_imag 

110 

111 else: 

112 # from numba_progress import ProgressBar 

113 # num_iterations = jmax * jmax * wavelengths 

114 # progress = ProgressBar(total=num_iterations) 

115 progress = None 

116 c_sca = compute_scattering_cross_section( 

117 lmax, 

118 particle_number, 

119 idx_lookup, 

120 f, 

121 translation_table, 

122 associated_legendre_lookup, 

123 spherical_bessel_lookup, 

124 e_j_dm_phi_loopup, 

125 progress, 

126 ) 

127 

128 self.c_sca = ( 

129 c_sca / np.power(np.abs(self.simulation.parameters.k_medium), 2) * np.pi 

130 ) 

131 

132 self.c_ext = np.real(self.c_ext) 

133 self.c_sca = np.real(self.c_sca) 

134 

135 self.albedo = self.c_sca / self.c_ext 

136 

137 def compute_phase_funcition( 

138 self, 

139 legendre_coefficients_number: int = 15, 

140 c_and_b: Union[bool, tuple] = False, 

141 ): 

142 """The function `compute_phase_function` calculates various polarization components and phase 

143 function coefficients for a given simulation. 

144 

145 Args: 

146 legendre_coefficients_number (int, optional): The number of Legendre coefficients to compute 

147 for the phase function. These coefficients are used to approximate the phase function 

148 using Legendre polynomials. The higher the number of coefficients, the more accurate 

149 the approximation will be. 

150 c_and_b (Union[bool, tuple], optional): A boolean value or a tuple. If `True`, it indicates 

151 that the `c` and `b` bounds should be computed. If `False`, the function will return 

152 without computing the `c` and `b` bounds. 

153 """ 

154 pilm, taulm = spherical_functions_trigon( 

155 self.simulation.numerics.lmax, self.simulation.numerics.polar_angles 

156 ) 

157 

158 if self.simulation.numerics.gpu: 

159 jmax = ( 

160 self.simulation.parameters.particles.number 

161 * self.simulation.numerics.nmax 

162 ) 

163 angles = self.simulation.numerics.azimuthal_angles.size 

164 wavelengths = self.simulation.parameters.k_medium.size 

165 e_field_theta_real = np.zeros( 

166 ( 

167 self.simulation.numerics.azimuthal_angles.size, 

168 self.simulation.parameters.k_medium.size, 

169 ), 

170 dtype=float, 

171 ) 

172 e_field_theta_imag = np.zeros_like(e_field_theta_real) 

173 e_field_phi_real = np.zeros_like(e_field_theta_real) 

174 e_field_phi_imag = np.zeros_like(e_field_theta_real) 

175 

176 particles_position_device = cuda.to_device( 

177 self.simulation.parameters.particles.position 

178 ) 

179 idx_device = cuda.to_device(self.simulation.idx_lookup) 

180 sfc_device = cuda.to_device(self.simulation.scattered_field_coefficients) 

181 k_medium_device = cuda.to_device(self.simulation.parameters.k_medium) 

182 azimuthal_angles_device = cuda.to_device( 

183 self.simulation.numerics.azimuthal_angles 

184 ) 

185 e_r_device = cuda.to_device(self.simulation.numerics.e_r) 

186 pilm_device = cuda.to_device(pilm) 

187 taulm_device = cuda.to_device(taulm) 

188 e_field_theta_real_device = cuda.to_device(e_field_theta_real) 

189 e_field_theta_imag_device = cuda.to_device(e_field_theta_imag) 

190 e_field_phi_real_device = cuda.to_device(e_field_phi_real) 

191 e_field_phi_imag_device = cuda.to_device(e_field_phi_imag) 

192 

193 sizes = (jmax, angles, wavelengths) 

194 threads_per_block = (16, 16, 2) 

195 # blocks_per_grid = tuple( 

196 # [ 

197 # ceil(sizes[k] / threads_per_block[k]) 

198 # for k in range(len(threads_per_block)) 

199 # ] 

200 # ) 

201 blocks_per_grid = tuple( 

202 ceil(sizes[k] / threads_per_block[k]) 

203 for k in range(len(threads_per_block)) 

204 ) 

205 # blocks_per_grid = ( 

206 # ceil(jmax / threads_per_block[0]), 

207 # ceil(angles / threads_per_block[1]), 

208 # ceil(wavelengths / threads_per_block[2])) 

209 

210 compute_electric_field_angle_components_gpu[ 

211 blocks_per_grid, threads_per_block 

212 ]( 

213 self.simulation.numerics.lmax, 

214 particles_position_device, 

215 idx_device, 

216 sfc_device, 

217 k_medium_device, 

218 azimuthal_angles_device, 

219 e_r_device, 

220 pilm_device, 

221 taulm_device, 

222 e_field_theta_real_device, 

223 e_field_theta_imag_device, 

224 e_field_phi_real_device, 

225 e_field_phi_imag_device, 

226 ) 

227 

228 e_field_theta_real = e_field_theta_real_device.copy_to_host() 

229 e_field_theta_imag = e_field_theta_imag_device.copy_to_host() 

230 e_field_phi_real = e_field_phi_real_device.copy_to_host() 

231 e_field_phi_imag = e_field_phi_imag_device.copy_to_host() 

232 e_field_theta = e_field_theta_real + 1j * e_field_theta_imag 

233 e_field_phi = e_field_phi_real + 1j * e_field_phi_imag 

234 

235 intensity = np.zeros_like(e_field_theta_real) 

236 dop = np.zeros_like(e_field_theta_real) 

237 dolp = np.zeros_like(e_field_theta_real) 

238 dolq = np.zeros_like(e_field_theta_real) 

239 dolu = np.zeros_like(e_field_theta_real) 

240 docp = np.zeros_like(e_field_theta_real) 

241 

242 intensity_device = cuda.to_device(intensity) 

243 dop_device = cuda.to_device(dop) 

244 dolp_device = cuda.to_device(dolp) 

245 dolq_device = cuda.to_device(dolq) 

246 dolu_device = cuda.to_device(dolu) 

247 docp_device = cuda.to_device(docp) 

248 

249 sizes = (angles, wavelengths) 

250 threads_per_block = (32, 32) 

251 # blocks_per_grid = tuple( 

252 # [ 

253 # ceil(sizes[k] / threads_per_block[k]) 

254 # for k in range(len(threads_per_block)) 

255 # ] 

256 # ) 

257 blocks_per_grid = tuple( 

258 ceil(sizes[k] / threads_per_block[k]) 

259 for k in range(len(threads_per_block)) 

260 ) 

261 compute_polarization_components_gpu[blocks_per_grid, threads_per_block]( 

262 self.simulation.parameters.k_medium.size, 

263 self.simulation.numerics.azimuthal_angles.size, 

264 e_field_theta_real_device, 

265 e_field_theta_imag_device, 

266 e_field_phi_real_device, 

267 e_field_phi_imag_device, 

268 intensity_device, 

269 dop_device, 

270 dolp_device, 

271 dolq_device, 

272 dolu_device, 

273 docp_device, 

274 ) 

275 

276 intensity = intensity_device.copy_to_host() 

277 dop = dop_device.copy_to_host() 

278 dolp = dolp_device.copy_to_host() 

279 dolq = dolq_device.copy_to_host() 

280 dolu = dolu_device.copy_to_host() 

281 docp = docp_device.copy_to_host() 

282 else: 

283 e_field_theta, e_field_phi = compute_electric_field_angle_components( 

284 self.simulation.numerics.lmax, 

285 self.simulation.parameters.particles.position, 

286 self.simulation.idx_lookup, 

287 self.simulation.scattered_field_coefficients, 

288 self.simulation.parameters.k_medium, 

289 self.simulation.numerics.azimuthal_angles, 

290 self.simulation.numerics.e_r, 

291 pilm, 

292 taulm, 

293 ) 

294 

295 intensity, dop, dolp, dolq, dolu, docp = compute_polarization_components( 

296 self.simulation.parameters.k_medium.size, 

297 self.simulation.numerics.azimuthal_angles.size, 

298 e_field_theta, 

299 e_field_phi, 

300 ) 

301 

302 self.scattering_angles = self.simulation.numerics.polar_angles 

303 

304 self.phase_function_3d = ( 

305 intensity 

306 * 4 

307 * np.pi 

308 / np.power(np.abs(self.simulation.parameters.k_medium), 2) 

309 / self.c_sca[np.newaxis, :] 

310 ) 

311 self.phase_function_legendre_coefficients = np.polynomial.legendre.legfit( 

312 np.cos(self.scattering_angles), 

313 self.phase_function_3d, 

314 legendre_coefficients_number, 

315 ) 

316 

317 self.degree_of_polarization_3d = dop 

318 self.degree_of_linear_polarization_3d = dolp 

319 self.degree_of_linear_polarization_q_3d = dolq 

320 self.degree_of_linear_polarization_u_3d = dolu 

321 self.degree_of_circular_polarization_3d = docp 

322 

323 if (self.simulation.numerics.sampling_points_number is not None) and ( 

324 self.simulation.numerics.sampling_points_number.size == 2 

325 ): 

326 self.phase_function = np.mean( 

327 np.reshape( 

328 self.phase_function_3d, 

329 np.append( 

330 self.simulation.numerics.sampling_points_number, 

331 self.simulation.parameters.k_medium.size, 

332 ), 

333 ), 

334 axis=0, 

335 ) 

336 

337 self.degree_of_polarization = np.mean( 

338 np.reshape( 

339 dop, 

340 np.append( 

341 self.simulation.numerics.sampling_points_number, 

342 self.simulation.parameters.k_medium.size, 

343 ), 

344 ), 

345 axis=0, 

346 ) 

347 self.degree_of_linear_polarization = np.mean( 

348 np.reshape( 

349 dolp, 

350 np.append( 

351 self.simulation.numerics.sampling_points_number, 

352 self.simulation.parameters.k_medium.size, 

353 ), 

354 ), 

355 axis=0, 

356 ) 

357 self.degree_of_linear_polarization_q = np.mean( 

358 np.reshape( 

359 dolq, 

360 np.append( 

361 self.simulation.numerics.sampling_points_number, 

362 self.simulation.parameters.k_medium.size, 

363 ), 

364 ), 

365 axis=0, 

366 ) 

367 self.degree_of_linear_polarization_u = np.mean( 

368 np.reshape( 

369 dolu, 

370 np.append( 

371 self.simulation.numerics.sampling_points_number, 

372 self.simulation.parameters.k_medium.size, 

373 ), 

374 ), 

375 axis=0, 

376 ) 

377 self.degree_of_circular_polarization = np.mean( 

378 np.reshape( 

379 docp, 

380 np.append( 

381 self.simulation.numerics.sampling_points_number, 

382 self.simulation.parameters.k_medium.size, 

383 ), 

384 ), 

385 axis=0, 

386 ) 

387 

388 self.scattering_angles = np.reshape( 

389 self.scattering_angles, self.simulation.numerics.sampling_points_number 

390 ) 

391 self.scattering_angles = self.scattering_angles[0, :] 

392 else: 

393 self.phase_function = self.phase_function_3d 

394 

395 self.degree_of_polarization = dop 

396 self.degree_of_linear_polarization = dolp 

397 self.degree_of_linear_polarization_q = dolq 

398 self.degree_of_linear_polarization_u = dolu 

399 self.degree_of_circular_polarization = docp 

400 

401 self.c_and_b_bounds = c_and_b 

402 if isinstance(c_and_b, bool): 

403 if c_and_b: 

404 self.c_and_b_bounds = ([-1, 0], [1, 1]) 

405 else: 

406 return 

407 

408 self.__compute_c_and_b() 

409 

410 @staticmethod 

411 def compute_double_henyey_greenstein(theta: np.ndarray, cb: np.ndarray): 

412 """ 

413 Calculates the scattering phase function using the double Henyey-Greenstein model. 

414 

415 Args: 

416 theta (np.ndarray): A numpy array representing the scattering angle. It contains the 

417 values at which you want to compute the double Henyey-Greenstein phase function. 

418 cb (np.ndarray): A numpy array that represents the coefficients of the double 

419 Henyey-Greenstein phase function. It should have a size of either 2 or 3. 

420 

421 Returns: 

422 p (np.ndarray): Computed Henyey-Greenstein phase function. 

423 """ 

424 cb = np.squeeze(cb) 

425 if cb.size < 2: 

426 cb = np.array([0, 0.5, 0.5]) 

427 elif cb.size == 2: 

428 cb = np.append(cb, cb[1]) 

429 elif cb.size > 3: 

430 cb = cb[:2] 

431 

432 p1 = (1 - cb[1] ** 2) / np.power( 

433 1 - 2 * cb[1] * np.cos(theta) + cb[1] ** 2, 3 / 2 

434 ) 

435 p2 = (1 - cb[2] ** 2) / np.power( 

436 1 + 2 * cb[2] * np.cos(theta) + cb[2] ** 2, 3 / 2 

437 ) 

438 return (1 - cb[0]) / 2 * p1 + (1 + cb[0]) / 2 * p2 

439 

440 def __compute_c_and_b(self): 

441 """The function computes the values of parameters 'b' and 'c' using the double Henyey-Greenstein 

442 phase function. 

443 """ 

444 # double henyey greenstein 

445 if len(self.c_and_b_bounds[0]) not in [2, 3]: 

446 self.c_and_b_bounds = ([-1, 0], [1, 1]) 

447 self.log.warning( 

448 "Number of parameters need to be 2 (b,c) or 3 (b1,b2,c). Reverting to two parameters (b,c) and setting the bounds to standard: b in [0, 1] and c in [-1, 1]" 

449 ) 

450 

451 from scipy.optimize import least_squares 

452 

453 if len(self.c_and_b_bounds) == 2: 

454 bc0 = np.array([0, 0.5]) 

455 else: 

456 bc0 = np.array([0, 0.5, 0.5]) 

457 

458 self.cb = np.empty((self.phase_function.shape[1], len(self.c_and_b_bounds))) 

459 for w in range(self.phase_function.shape[1]): 

460 

461 def dhg_optimization(bc): 

462 return ( 

463 Optics.compute_double_henyey_greenstein(self.scattering_angles, bc) 

464 - self.phase_function[:, w] 

465 ) 

466 

467 bc = least_squares( 

468 dhg_optimization, bc0, jac="2-point", bounds=self.c_and_b_bounds 

469 ) 

470 self.cb[w, :] = bc.x