Coverage for yasfpy/simulation.py: 72%

1import yasfpy.log as log

2from time import time

4import numpy as np

5from math import ceil

6from numba import cuda

7from scipy.sparse.linalg import LinearOperator

9# from scipy.spatial.distance import pdist, squareform

10from scipy.special import spherical_jn, spherical_yn

12# from scipy.special import hankel1

13# from scipy.special import lpmv

15from yasfpy.parameters import Parameters

16from yasfpy.numerics import Numerics

17from yasfpy.functions.spherical_functions_trigon import spherical_functions_trigon

18from yasfpy.functions.t_entry import t_entry

20from yasfpy.functions.misc import transformation_coefficients

21from yasfpy.functions.misc import multi2single_index

22from yasfpy.functions.misc import mutual_lookup

24from yasfpy.functions.cpu_numba import compute_idx_lookups

25from yasfpy.functions.cpu_numba import particle_interaction, compute_field

26from yasfpy.functions.cuda_numba import particle_interaction_gpu, compute_field_gpu

29class Simulation:

30 """This class represents the simulation of YASF (Yet Another Scattering Framework).

31 It contains methods for initializing the simulation, computing lookup tables, and calculating mie coefficients.

32 """

34 def __init__(self, parameters: Parameters, numerics: Numerics):

35 """

36 Initialize the Simulation object.

38 Args:

39 parameters (Parameters): The parameters for the simulation.

40 numerics (Numerics): The numerics for the simulation.

41 """

42 self.parameters = parameters

43 self.numerics = numerics

45 self.log = log.scattering_logger(__name__)

46 self.__setup()

48 def legacy_compute_lookup_particle_distances(self):

49 """

50 The largest distance between two particles is divided into segments provided by `Numerics.particle_distance_resolution`.

51 This array is then used as a lookup for the calculation of the spherical Hankel function.

53 Notes

54 -----

55 This function has been ported from the Matlab Celes framework but is not used by YASF!

56 """

57 # add two zeros at the beginning to allow interpolation

58 # also in the first segment

59 step = self.numerics.particle_distance_resolution

60 maxdist = (

61 self.parameters.particles.max_particle_distance

62 + 3 * self.numerics.particle_distance_resolution

63 )

64 self.lookup_particle_distances = np.concatenate(

65 (np.array([0]), np.arange(0, maxdist + np.finfo(float).eps, step))

66 )

68 def legacy_compute_h3_table(self):

69 """

70 Computes the spherical hankel function

71 at the points calculated in `Simulation.legacy_compute_lookup_particle_distances()`.

73 Attributes:

74 h3_table (np.ndarray): Lookup table of the spherical hankel function values at `self.lookup_particle_distances`

76 Notes:

77 This function has been ported from the Matlab Celes framework but is not used by YASF!

78 """

79 self.h3_table = np.zeros(

80 (

81 2 * self.numerics.lmax + 1,

82 self.lookup_particle_distances.shape[0],

83 self.parameters.medium_refractive_index.shape[0],

84 ),

85 dtype=complex,

86 )

87 size_param = np.outer(self.lookup_particle_distances, self.parameters.k_medium)

89 for p in range(2 * self.numerics.lmax + 1):

90 self.h3_table[p, :, :] = spherical_jn(p, size_param) + 1j * spherical_yn(

91 p, size_param

92 )

94 def __compute_idx_lookup(self):

95 """

96 Creates a lookup table with the indices used in further calculations.

97 The lookup table is created using `compute_idx_lookups` function from `yasfpy.functions.cpu_numba`.

99 Attributes:

100 idx_lookup (np.ndarray): Lookup table of the indices to iterate over large arrays.

101

102 Notes:

103 This function utilizes Numba to optimize the computations.

104 """

105 self.idx_lookup = compute_idx_lookups(

106 self.numerics.lmax, self.parameters.particles.number

107 )

108

109 def __compute_lookups(self):

110 """

111 Computes various lookup tables for each particle.

112

113 Attributes:

114 sph_j (np.ndarray): Spherical Bessel function lookup table calculated for pair-wise particle distances.

115 sph_h (np.ndarray): Spherical Hankel function lookup table calculated for pair-wise particle distances.

116 plm (np.ndarray): Associated Legendre polynomial lookup table calculated for the cosine value of the pairwise particle inclination angles.

117 e_j_dm_phi (np.ndarray): Exponential function lookup table calculated for the pairwise particle azimuthal angles.

118

119 Notes:

120 This function uses numba (https://numba.pydata.org/) under the hood to speed up the computations.

121 """

122 lookup_computation_time_start = time()

123 # TODO: new, could be error prone and is not tested yet!

124 self.sph_j, self.sph_h, self.e_j_dm_phi, self.plm = mutual_lookup(

125 self.numerics.lmax,

126 self.parameters.particles.position,

127 self.parameters.particles.position,

128 self.parameters.k_medium,

129 )[:4]

130

131 # lmax = self.numerics.lmax

132 # particle_number = self.parameters.particles.number

133

134 # dists = squareform(pdist(self.parameters.particles.position))

135 # ct = np.divide(

136 # np.subtract.outer(

137 # self.parameters.particles.position[:, 2], self.parameters.particles.position[:, 2]),

138 # dists,

139 # out = np.zeros((particle_number, particle_number)),

140 # where = dists != 0)

141 # phi = np.arctan2(

142 # np.subtract.outer(

143 # self.parameters.particles.position[:, 1], self.parameters.particles.position[:, 1]),

144 # np.subtract.outer(self.parameters.particles.position[:, 0], self.parameters.particles.position[:, 0]))

145

146 # size_param = np.outer(dists.ravel(), self.parameters.k_medium).reshape(

147 # [particle_number, particle_number, self.parameters.k_medium.shape[0]])

148

149 # self.sph_h = np.zeros((2 * lmax + 1, particle_number, particle_number, self.parameters.k_medium.shape[0]), dtype=complex)

150 # self.sph_j = np.zeros_like(self.sph_h)

151 # self.e_j_dm_phi = np.zeros((4 * lmax + 1, particle_number, particle_number), dtype=complex)

152 # self.plm = np.zeros(((lmax + 1) * (2 * lmax + 1),

153 # particle_number, particle_number))

154

155 # for p in range(2 * lmax + 1):

156 # self.sph_h[p, :, :, :] = np.sqrt(

157 # np.divide(

158 # np.pi / 2,

159 # size_param,

160 # out=np.zeros_like(size_param),

161 # where=size_param != 0)

162 # ) * hankel1(p + 1/2, size_param)

163 # self.sph_j[p, :, :, :] = spherical_jn(p, size_param)

164 # self.e_j_dm_phi[p, :, :] = np.exp(1j * (p - 2 * lmax) * phi)

165 # self.e_j_dm_phi[p + 2 * lmax, :, :] = np.exp(1j * p * phi)

166 # for absdm in range(p + 1):

167 # cml = np.sqrt((2 * p + 1) / 2 /

168 # np.prod(np.arange(p - absdm + 1, p + absdm + 1)))

169 # self.plm[p * (p + 1) // 2 + absdm, :, :] = cml * \

170 # np.power(-1.0, absdm) * lpmv(absdm, p, ct)

171

172 # self.sph_h = np.nan_to_num(

173 # self.sph_h, nan=0) + np.isnan(self.sph_h) * 1

174

175 lookup_computation_time_stop = time()

176 self.log.scatter(

177 "Computing lookup tables took %f s"

178 % (lookup_computation_time_stop - lookup_computation_time_start)

179 )

180

181 def __setup(self):

182 """

183 An internal setup function called upon object creation.

184 The following functions are called:

185

186 - [__compute_idx_lookups][yasfpy.simulation.Simulation.__compute_idx_lookup]

187 - [__compute_lookups][yasfpy.simulation.Simulation.__compute_lookups]

188 """

189 self.__compute_idx_lookup()

190 self.__compute_lookups()

191

192 def compute_mie_coefficients(self):

193 """

194 Computes the mie coefficients for the unique pair

195 of particle radius and the refractive index of the particle.

196

197 Attributes:

198 mie_coefficients (np.ndarray): Mie coefficients table

199

200 See Also:

201 [t_entry][yasfpy.functions.t_entry.t_entry] : T-Matrix entry function

202

203 Notes:

204 Due to the four nested loops (particles, tau, l, and m),

205 it could be rewritten using `numba` to speed the process up.

206 """

207 self.mie_coefficients = np.zeros(

208 (

209 self.parameters.particles.num_unique_pairs,

210 self.numerics.nmax,

211 self.parameters.wavelength.shape[0],

212 ),

213 dtype=complex,

214 )

215

216 self.scatter_to_internal = np.zeros_like(self.mie_coefficients)

217

218 for u_i in range(self.parameters.particles.num_unique_pairs):

219 for tau in range(1, 3):

220 for l in range(1, self.numerics.lmax + 1):

221 for m in range(-l, l + 1):

222 jmult = multi2single_index(0, tau, l, m, self.numerics.lmax)

223 self.mie_coefficients[u_i, jmult, :] = t_entry(

224 tau=tau,

225 l=l,

226 k_medium=self.parameters.k_medium,

227 k_sphere=self.parameters.omega

228 * self.parameters.particles.unique_radius_index_pairs[

229 u_i, 1:

230 ],

231 radius=np.real(

232 self.parameters.particles.unique_radius_index_pairs[

233 u_i, 0

234 ]

235 ),

236 )

237

238 self.scatter_to_internal[u_i, jmult, :] = t_entry(

239 tau=tau,

240 l=l,

241 k_medium=self.parameters.k_medium,

242 k_sphere=self.parameters.omega

243 * self.parameters.particles.unique_radius_index_pairs[

244 u_i, 1:

245 ],

246 radius=np.real(

247 self.parameters.particles.unique_radius_index_pairs[

248 u_i, 0

249 ]

250 ),

251 field_type="ratio",

252 )

253

254 def compute_initial_field_coefficients(self):

255 r"""

256 Computes initial field coefficients $a_{\\tau ,l,m}$ and $b_{\\tau ,l,m}$.

257 Depending on the `beam_width`, one of two functions is called:

258

259 - [__compute_initial_field_coefficients_wavebundle_normal_incidence][yasfpy.simulation.Simulation.__compute_initial_field_coefficients_wavebundle_normal_incidence], $\\text{beam width} \\in (0, \\infty)$

260 - [__compute_initial_field_coefficients_planewave][yasfpy.simulation.Simulation.__compute_initial_field_coefficients_planewave], $\\text{beam width} = 0$ or $\\text{beam width} = \\infty$

261

262 Attributes:

263 initial_field_coefficients (np.ndarray): Initial field coefficients

264 """

265 self.log.scatter("compute initial field coefficients ...")

266

267 if np.isfinite(self.parameters.initial_field.beam_width) and (

268 self.parameters.initial_field.beam_width > 0

269 ):

270 self.log.scatter("\t Gaussian beam ...")

271 if self.parameters.initial_field.normal_incidence:

272 self.__compute_initial_field_coefficients_wavebundle_normal_incidence()

273 else:

274 self.log.error("\t this case is not implemented")

275 else:

276 self.log.scatter("\t plane wave ...")

277 self.__compute_initial_field_coefficients_planewave()

278

279 self.log.scatter("done")

280

281 def compute_right_hand_side(self):

282 r"""

283 Computes the right hand side $T \\cdot a_I$ of the equation $M \\cdot b = T \\cdot a_I$.

284

285 Attributes

286 ----------

287 right_hand_side : np.ndarray

288 Right hand side of the equation $M \\cdot b = T \\cdot a_I$

289

290 Notes

291 -----

292 For more information regarding the equation, please refer to the paper by Celes (https://arxiv.org/abs/1706.02145).

293 """

294 self.right_hand_side = (

295 self.mie_coefficients[self.parameters.particles.single_unique_array_idx, :]

296 * self.initial_field_coefficients

297 )

298

299 def __compute_initial_field_coefficients_planewave(self):

300 """The function computes the initial field coefficients for a plane wave based on given parameters

301 and spherical coordinates.

302

303 """

304 lmax = self.numerics.lmax

305 E0 = self.parameters.initial_field.amplitude

306 k = self.parameters.k_medium

307

308 beta = self.parameters.initial_field.polar_angle

309 cb = np.cos(beta)

310 sb = np.sin(beta)

311 alpha = self.parameters.initial_field.azimuthal_angle

312

313 # pi and tau symbols for transformation matrix B_dagger

314 pilm, taulm = spherical_functions_trigon(lmax, beta)

315

316 # cylindrical coordinates for relative particle positions

317 relative_particle_positions = (

318 self.parameters.particles.position

319 - self.parameters.initial_field.focal_point

320 )

321 kvec = np.outer(np.array((sb * np.cos(alpha), sb * np.sin(alpha), cb)), k)

322 eikr = np.exp(1j * np.matmul(relative_particle_positions, kvec))

323

324 # clean up some memory?

325 del (k, beta, cb, sb, kvec, relative_particle_positions)

326

327 self.initial_field_coefficients = np.zeros(

328 (

329 self.parameters.particles.number,

330 self.numerics.nmax,

331 self.parameters.k_medium.size,

332 ),

333 dtype=complex,

334 )

335 for m in range(-lmax, lmax + 1):

336 for tau in range(1, 3):

337 for l in range(np.max([1, np.abs(m)]), lmax + 1):

338 n = multi2single_index(0, tau, l, m, lmax)

339 self.initial_field_coefficients[:, n, :] = (

340 4

341 * E0

342 * np.exp(-1j * m * alpha)

343 * eikr

344 * transformation_coefficients(

345 pilm,

346 taulm,

347 tau,

348 l,

349 m,

350 self.parameters.initial_field.pol,

351 dagger=True,

352 )

353 )

354

355 def __compute_initial_field_coefficients_wavebundle_normal_incidence(self):

356 """The function initializes the field coefficients for a wave bundle incident at normal incidence.

357

358 TODO:

359 Implement this function using the celes function [initial_field_coefficients_wavebundle_normal_incidence.m](https://github.com/disordered-photonics/celes/blob/master/src/initial/initial_field_coefficients_wavebundle_normal_incidence.m)

360 """

361 self.initial_field_coefficients = (

362 np.zeros(

363 (

364 self.parameters.particles.number,

365 self.numerics.nmax,

366 self.parameters.k_medium.size,

367 ),

368 dtype=complex,

369 )

370 * np.nan

371 )

372

373 def coupling_matrix_multiply(self, x: np.ndarray, idx: int = None):

374 """Computes the coupling matrix `wx` based on the input parameters.

375

376 Args:

377 x (np.ndarray): An input array of shape (n,) or (n, m), where n is the number of particles and m is the number

378 of features for each particle. This array represents the input data for which the coupling

379 matrix needs to be computed.

380 idx (int): An optional integer that specifies the index of a specific spherical harmonic mode. If `idx` is provided,

381 the computation will only be performed for that specific mode. If `idx` is not provided or set to `None`,

382 the computation will be performed for all spherical harmonic modes.

383

384 Returns:

385 wx (np.ndarray): An array of shape (n, m, p), where n is the number of particles, m is the number of features for each

386 particle, and p is the number of wavelengths. It represents the coupling matrix `wx`.

387 """

388 self.log.scatter("prepare particle coupling ... ")

389 preparation_time = time()

390

391 lmax = self.numerics.lmax

392 particle_number = self.parameters.particles.number

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

394 wavelengths_size = self.parameters.k_medium.shape[0]

395 translation_table = self.numerics.translation_ab5

396 associated_legendre_lookup = self.plm

397 spherical_hankel_lookup = self.sph_h

398 e_j_dm_phi_loopup = self.e_j_dm_phi

399

400 idx_lookup = self.idx_lookup

401

402 if idx != None:

403 spherical_hankel_lookup = spherical_hankel_lookup[:, :, :, idx]

404 spherical_hankel_lookup = np.copy(

405 spherical_hankel_lookup[:, :, :, np.newaxis]

406 )

407 wavelengths_size = 1

408

409 self.log.scatter("\t Starting Wx computation")

410 if self.numerics.gpu:

411 wx_real = np.zeros(x.shape + (wavelengths_size,), dtype=float)

412 wx_imag = np.zeros_like(wx_real)

413

414 idx_device = cuda.to_device(idx_lookup)

415 x_device = cuda.to_device(x)

416 wx_real_device = cuda.to_device(wx_real)

417 wx_imag_device = cuda.to_device(wx_imag)

418 translation_device = cuda.to_device(translation_table)

419 associated_legendre_device = cuda.to_device(associated_legendre_lookup)

420 spherical_hankel_device = cuda.to_device(spherical_hankel_lookup)

421 e_j_dm_phi_device = cuda.to_device(e_j_dm_phi_loopup)

422

423 threads_per_block = (16, 16, 2)

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

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

426 blocks_per_grid_z = ceil(wavelengths_size / threads_per_block[2])

427 blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y, blocks_per_grid_z)

428

429 coupling_matrix_time = time()

430 particle_interaction_gpu[blocks_per_grid, threads_per_block](

431 lmax,

432 particle_number,

433 idx_device,

434 x_device,

435 wx_real_device,

436 wx_imag_device,

437 translation_device,

438 associated_legendre_device,

439 spherical_hankel_device,

440 e_j_dm_phi_device,

441 )

442 wx_real = wx_real_device.copy_to_host()

443 wx_imag = wx_imag_device.copy_to_host()

444 wx = wx_real + 1j * wx_imag

445 # particle_interaction.parallel_diagnostics(level=4)

446 time_end = time()

447 self.log.scatter(

448 "\t Time taken for preparation: %f"

449 % (coupling_matrix_time - preparation_time)

450 )

451 self.log.scatter(

452 "\t Time taken for coupling matrix: %f"

453 % (time_end - coupling_matrix_time)

454 )

455 else:

456 # from numba_progress import ProgressBar

457 # num_iterations = jmax * jmax * wavelengths

458 # progress = ProgressBar(total=num_iterations)

459 # progress = None

460 wx = particle_interaction(

461 lmax,

462 particle_number,

463 idx_lookup,

464 x,

465 translation_table,

466 associated_legendre_lookup,

467 spherical_hankel_lookup,

468 e_j_dm_phi_loopup,

469 )

470 time_end = time()

471 self.log.scatter(

472 "\t Time taken for coupling matrix: %f" % (time_end - preparation_time)

473 )

474

475 if idx != None:

476 wx = np.squeeze(wx)

477

478 return wx

479

480 def master_matrix_multiply(self, value: np.ndarray, idx: int):

481 """Applies a T-matrix to a given value and returns the result.

482

483 Args:

484 value (np.ndarray): The input value for the matrix multiplication operation.

485 idx (int): The index of the matrix to be multiplied.

486

487 Returns:

488 mx (np.ndarray): The result of the matrix multiplication operation.

489

490 """

491 wx = self.coupling_matrix_multiply(value, idx)

492

493 self.log.scatter("apply T-matrix ...")

494 t_matrix_start = time()

495

496 twx = (

497 self.mie_coefficients[

498 self.parameters.particles.single_unique_array_idx, :, idx

499 ].ravel(order="C")

500 * wx

501 )

502 mx = value - twx

503

504 t_matrix_stop = time()

505 self.log.scatter(f"\t done in {t_matrix_stop - t_matrix_start} seconds.")

506

507 return mx

508

509 def compute_scattered_field_coefficients(self, guess: np.ndarray = None):

510 """The function computes the scattered field coefficients using a linear operator and a solver.

511

512 Args:

513 guess (np.ndarray): Optional. The initial guess for the solution of the linear system. If no guess is provided,

514 the `right_hand_side` variable is used as the initial guess.

515

516 """

517 self.log.scatter("compute scattered field coefficients ...")

518 jmax = self.parameters.particles.number * self.numerics.nmax

519 self.scattered_field_coefficients = np.zeros_like(

520 self.initial_field_coefficients

521 )

522 self.scattered_field_err_codes = np.zeros(self.parameters.wavelengths_number)

523 if guess is None:

524 guess = self.right_hand_side

525 for w in range(self.parameters.wavelengths_number):

526

527 def mmm(x):

528 return self.master_matrix_multiply(x, w)

529

530 A = LinearOperator(shape=(jmax, jmax), matvec=mmm)

531 b = self.right_hand_side[:, :, w].ravel()

532 x0 = guess[:, :, w].ravel()

533 self.log.scatter(

534 "Solver run %d/%d" % (w + 1, self.parameters.wavelengths_number)

535 )

536 x, err_code = self.numerics.solver.run(A, b, x0)

537 self.scattered_field_coefficients[:, :, w] = x.reshape(

538 self.right_hand_side.shape[:2]

539 )

540 self.scattered_field_err_codes[w] = err_code

541

542 def compute_fields(self, sampling_points: np.ndarray):

543 """The function `compute_fields` calculates the field at given sampling points using either CPU or

544 GPU computation.

545

546 Args:

547 sampling_points (np.ndarray): The numpy array that represents the coordinates of the sampling points.

548 It should have a shape of `(n, 3)`, where `n` is the number of sampling points and each row

549 represents the `(x, y, z)` coordinates of a point.

550 """

551 if sampling_points.shape[0] < 1:

552 self.log.error("Number of sampling points must be bigger than zero!")

553 return

554 elif sampling_points.shape[1] != 3:

555 self.log.error("The points have to have three coordinates (x,y,z)!")

556 return

557

558 # scatter_to_internal_table = np.sum((self.parameters.particles.position[:, np.newaxis, :] - sampling_points[np.newaxis, :, :])**2, axis = 2)

559 # scatter_to_internal_table = scatter_to_internal_table < self.parameters.particles.r[:, np.newaxis]**2

560

561 print("Computing mutual lookup")

562 (

563 _,

564 sph_h,

565 e_j_dm_phi,

566 p_lm,

567 e_r,

568 e_theta,

569 e_phi,

570 cosine_theta,

571 sine_theta,

572 size_parameter,

573 sph_h_derivative,

574 ) = mutual_lookup(

575 self.numerics.lmax,

576 self.parameters.particles.position,

577 sampling_points,

578 self.parameters.k_medium,

579 derivatives=True,

580 parallel=False,

581 )

582 pi_lm, tau_lm = spherical_functions_trigon(

583 self.numerics.lmax, cosine_theta, sine_theta

584 )

585 # print(sph_h.size)

586

587 print("Computing field...")

588 field_time_start = time()

589 self.sampling_points = sampling_points

590 if self.numerics.gpu:

591 print("\t...using GPU")

592 field_real = np.zeros(

593 (self.parameters.k_medium.size, sampling_points.shape[0], 3),

594 dtype=float,

595 )

596 field_imag = np.zeros_like(field_real)

597

598 idx_device = cuda.to_device(self.idx_lookup)

599 size_parameter_device = cuda.to_device(np.ascontiguousarray(size_parameter))

600 sph_h_device = cuda.to_device(np.ascontiguousarray(sph_h))

601 sph_h_derivative_device = cuda.to_device(

602 np.ascontiguousarray(sph_h_derivative)

603 )

604 e_j_dm_phi_device = cuda.to_device(np.ascontiguousarray(e_j_dm_phi))

605 p_lm_device = cuda.to_device(np.ascontiguousarray(p_lm))

606 pi_lm_device = cuda.to_device(np.ascontiguousarray(pi_lm))

607 tau_lm_device = cuda.to_device(np.ascontiguousarray(tau_lm))

608 e_r_device = cuda.to_device(np.ascontiguousarray(e_r))

609 e_theta_device = cuda.to_device(np.ascontiguousarray(e_theta))

610 e_phi_device = cuda.to_device(np.ascontiguousarray(e_phi))

611 sfc_device = cuda.to_device(

612 np.ascontiguousarray(self.scattered_field_coefficients)

613 )

614

615 field_real_device = cuda.to_device(field_real)

616 field_imag_device = cuda.to_device(field_imag)

617

618 threads_per_block = (16, 16, 2)

619 blocks_per_grid = (

620 sampling_points.shape[0],

621 sph_h.shape[1] * 2 * self.numerics.lmax * (self.numerics.lmax + 2),

622 self.parameters.k_medium.size,

623 )

624 # blocks_per_grid = tuple(

625 # [

626 # ceil(blocks_per_grid[i] / threads_per_block[i])

627 # for i in range(len(threads_per_block))

628 # ]

629 # )

630 blocks_per_grid = tuple(

631 ceil(blocks_per_grid[i] / threads_per_block[i])

632 for i in range(len(threads_per_block))

633 )

634

635 compute_field_gpu[blocks_per_grid, threads_per_block](

636 self.numerics.lmax,

637 idx_device,

638 size_parameter_device,

639 sph_h_device,

640 sph_h_derivative_device,

641 e_j_dm_phi_device,

642 p_lm_device,

643 pi_lm_device,

644 tau_lm_device,

645 e_r_device,

646 e_theta_device,

647 e_phi_device,

648 sfc_device,

649 field_real_device,

650 field_imag_device,

651 )

652

653 field_real = field_real_device.copy_to_host()

654 field_imag = field_imag_device.copy_to_host()

655 self.scattered_field = field_real + 1j * field_imag

656

657 else:

658 print("\t...using CPU")

659 self.scattered_field = compute_field(

660 self.numerics.lmax,

661 self.idx_lookup,

662 size_parameter,

663 sph_h,

664 sph_h_derivative,

665 e_j_dm_phi,

666 p_lm,

667 pi_lm,

668 tau_lm,

669 e_r,

670 e_theta,

671 e_phi,

672 scattered_field_coefficients=self.scattered_field_coefficients,

673 )

674

675 field_time_stop = time()

676 self.log.scatter(

677 f"\t Time taken for field calculation: {field_time_stop - field_time_start}"

678 )