Coverage for pySDC/implementations/problem_classes/PenningTrap_3D.py: 92%

120 statements

, created at 2024-09-20 17:10 +0000

1import numpy as np

2from numba import jit

4from pySDC.core.errors import ProblemError

5from pySDC.core.problem import Problem, WorkCounter

6from pySDC.implementations.datatype_classes.particles import particles, fields, acceleration

9# noinspection PyUnusedLocal

10class penningtrap(Problem):

11 r"""

12 This class implements a standard Penning trap problem on the time interval :math:[0, t_{end}]

13 fully investigated in [1]_. The equations are given by the following equation of motion

15 .. math::

16 \frac{dv}{dt}=f(x,v)=\alpha[E(x,t)+v\times B(x,t)],

18 .. math::

19 \frac{dx}{dt}=v

21 with the particles :math:x, v\in \mathbb{R}^{3}. For the penning trap problem, the other parameters are given by

22 the constant magnetic field :math:B=\frac{\omega_{B}}{\alpha}\cdot \hat{e_{z}}\in \mathbb{R}^{3}

23 along the :math:z-axis with the particle's charge-to-mass ratio :math:\alpha=\frac{q}{m} so that

25 .. math::

26 v\times B=\frac{\omega_{B}}{\alpha}\left(

27 \begin{matrix}

28 0 & 1 & 0\\

29 -1 & 0 & 0\\

30 0 & 0 & 0

31 \end{matrix}

32 \right)v.

34 The electric field :math:E(x_{i})=E_{ext}(x_{i})+E_{int}(x_{i})\in \mathbb{R}^{3} where

36 .. math::

37 E_{ext}(x_{i})=-\epsilon\frac{\omega_{E}^{2}}{\alpha}\left(

38 \begin{matrix}

39 1 & 0 & 0\\

40 0 & 1 & 0\\

41 0 & 0 & -2

42 \end{matrix}

43 \right)x_{i}

45 and the inter-particle Coulomb interaction

47 .. math::

48 E_{int}(x_{i})=\sum_{k=1, k\neq i}^{N_{particles}}Q_{k}\frac{x_{i}-x_{k}}{(|x_{i}-x_{k}|^{2}+\lambda^{2})^{3/2}}

50 with the smoothing parameter :math:\lambda>0.

51 The exact solution also given for the single particle penning trap more detailed [1]_, [2]_.

52 For to solve nonlinear equation of system, Boris trick is used (see [2]_).

54 Parameters

55 ----------

56 omega_B : float

57 Amplitude of magnetic field.

58 omega_E : float

59 Amplitude of electric field.

60 u0 : np.1darray

61 Initial condition for position, and for velocity.

62 q : np.1darray

63 Particle's charge.

64 m : np.1darray

65 Mass.

66 nparts : int

67 The number of particles.

68 sig : float

69 The smoothing parameter :math:\lambda>0.

71 Attributes

72 ----------

73 work_counter : dict

74 Counts the calls of the right-hand side, and calls of the Boris solver.

76 References

77 ----------

78 .. [1] F. Penning. Die Glimmentladung bei niedrigem Druck zwischen koaxialen Zylindern in einem axialen Magnetfeld.

79 Physica. Vol. 3 (1936).

80 .. [2] Mathias Winkel, Robert Speck and Daniel Ruprecht. A high-order Boris integrator.

81 Journal of Computational Physics (2015).

82 """

84 dtype_u = particles

85 dtype_f = fields

87 def __init__(self, omega_B, omega_E, u0, nparts, sig):

88 # invoke super init, passing nparts, dtype_u and dtype_f

89 super().__init__(((3, nparts), None, np.dtype('float64')))

91 self._makeAttributeAndRegister('omega_B', 'omega_E', 'u0', 'sig', localVars=locals())

92 self.work_counters['Boris_solver'] = WorkCounter()

93 self.work_counters['rhs'] = WorkCounter()

95 @staticmethod

96 @jit(nopython=True, nogil=True)

97 def fast_interactions(N, pos, sig, q):

98 r"""

99 Computes the fast interactions.

101 Parameters

102 ----------

103 N : int

104 Number of particles.

105 pos : np.2darray

106 Position.

107 sig : float

108 The smoothing parameter :math:\lambda > 0.

109 q : np.1darray

110 Particle's charge.

112 Returns

113 -------

114 Efield : np.2darray

115 The internal E field for each particle.

116 """

117 Efield = np.zeros((3, N))

118 contrib = np.zeros(3)

120 for i in range(N):

121 contrib[:] = 0

123 for j in range(N):

124 dist2 = (

125 (pos[0, i] - pos[0, j]) ** 2 + (pos[1, i] - pos[1, j]) ** 2 + (pos[2, i] - pos[2, j]) ** 2 + sig**2

126 )

127 contrib += q[j] * (pos[:, i] - pos[:, j]) / dist2**1.5

129 Efield[:, i] += contrib[:]

131 return Efield

133 def get_interactions(self, part):

134 r"""

135 Routine to compute the particle-particle interaction, assuming :math:q = 1 for all particles.

137 Parameters

138 ----------

139 part : dtype_u

140 The particles.

142 Returns

143 -------

144 Efield : np.ndarray

145 The internal E field for each particle.

146 """

148 N = self.nparts

150 Efield = self.fast_interactions(N, part.pos, self.sig, part.q)

152 return Efield

154 def eval_f(self, part, t):

155 """

156 Routine to compute the E and B fields (named f for consistency with the original PEPC version).

158 Parameters

159 ----------

160 part : dtype_u

161 The particles.

162 t : float

163 Current time of the particles (not used here).

165 Returns

166 -------

167 f : dtype_f

168 Fields for the particles (internal and external), i.e., the right-hand side of the problem.

169 """

171 N = self.nparts

173 self.work_counters['rhs']()

174 try:

175 _unused = penningtrap.Harmonic_oscillator

176 Emat = np.diag([0, 0, -1])

177 except AttributeError:

178 Emat = np.diag([1, 1, -2])

180 f = self.dtype_f(self.init)

182 f.elec[:] = self.get_interactions(part)

184 for n in range(N):

185 f.elec[:, n] += self.omega_E**2 / (part.q[n] / part.m[n]) * np.dot(Emat, part.pos[:, n])

186 f.magn[:, n] = self.omega_B * np.array([0, 0, 1])

188 return f

190 # TODO : Warning, this should be moved to u_exact(t=0) !

191 def u_init(self):

192 """

193 Routine to compute the starting values for the particles.

195 Returns

196 -------

197 u : dtype_u

198 Particle set filled with initial data.

199 """

201 u0 = self.u0

202 N = self.nparts

204 u = self.dtype_u(self.init)

206 if u0[2][0] != 1 or u0[3][0] != 1:

207 raise ProblemError('so far only q = m = 1 is implemented')

209 # set first particle to u0

210 u.pos[0, 0] = u0[0][0]

211 u.pos[1, 0] = u0[0][1]

212 u.pos[2, 0] = u0[0][2]

213 u.vel[0, 0] = u0[1][0]

214 u.vel[1, 0] = u0[1][1]

215 u.vel[2, 0] = u0[1][2]

217 u.q[0] = u0[2][0]

218 u.m[0] = u0[3][0]

220 # initialize random seed

221 np.random.seed(N)

223 comx = u.pos[0, 0]

224 comy = u.pos[1, 0]

225 comz = u.pos[2, 0]

227 for n in range(1, N):

228 # draw 3 random variables in [-1,1] to shift positions

229 r = np.random.random_sample(3) - 1

230 u.pos[0, n] = r[0] + u0[0][0]

231 u.pos[1, n] = r[1] + u0[0][1]

232 u.pos[2, n] = r[2] + u0[0][2]

234 # draw 3 random variables in [-5,5] to shift velocities

235 r = np.random.random_sample(3) - 5

236 u.vel[0, n] = r[0] + u0[1][0]

237 u.vel[1, n] = r[1] + u0[1][1]

238 u.vel[2, n] = r[2] + u0[1][2]

240 u.q[n] = u0[2][0]

241 u.m[n] = u0[3][0]

243 # gather positions to check center

244 comx += u.pos[0, n]

245 comy += u.pos[1, n]

246 comz += u.pos[2, n]

248 # print('Center of positions:',comx/N,comy/N,comz/N)

250 return u

252 def u_exact(self, t):

253 r"""

254 Routine to compute the exact trajectory at time :math:t (only for single-particle setup).

256 Parameters

257 ----------

258 t : float

259 Current time of the exact trajectory.

261 Returns

262 -------

263 u : dtype_u

264 Particle type containing the exact position and velocity.

265 """

267 # some abbreviations

268 wE = self.omega_E

269 wB = self.omega_B

270 N = self.nparts

271 u0 = self.u0

273 if N != 1:

274 raise ProblemError('u_exact is only valid for a single particle')

276 u = self.dtype_u(((3, 1), self.init[1], self.init[2]))

278 wbar = np.sqrt(2) * wE

280 # position and velocity in z direction is easy to compute

281 u.pos[2, 0] = u0[0][2] * np.cos(wbar * t) + u0[1][2] / wbar * np.sin(wbar * t)

282 u.vel[2, 0] = -u0[0][2] * wbar * np.sin(wbar * t) + u0[1][2] * np.cos(wbar * t)

284 # define temp. variables to compute complex position

285 Op = 1 / 2 * (wB + np.sqrt(wB**2 - 4 * wE**2))

286 Om = 1 / 2 * (wB - np.sqrt(wB**2 - 4 * wE**2))

287 Rm = (Op * u0[0][0] + u0[1][1]) / (Op - Om)

288 Rp = u0[0][0] - Rm

289 Im = (Op * u0[0][1] - u0[1][0]) / (Op - Om)

290 Ip = u0[0][1] - Im

292 # compute position in complex notation

293 w = (Rp + Ip * 1j) * np.exp(-Op * t * 1j) + (Rm + Im * 1j) * np.exp(-Om * t * 1j)

294 # compute velocity as time derivative of the position

295 dw = -1j * Op * (Rp + Ip * 1j) * np.exp(-Op * t * 1j) - 1j * Om * (Rm + Im * 1j) * np.exp(-Om * t * 1j)

297 # get the appropriate real and imaginary parts

298 u.pos[0, 0] = w.real

299 u.vel[0, 0] = dw.real

300 u.pos[1, 0] = w.imag

301 u.vel[1, 0] = dw.imag

303 return u

305 def build_f(self, f, part, t):

306 """

307 Helper function to assemble the correct right-hand side out of B and E field.

309 Parameters

310 ----------

311 f : dtype_f

312 The field values.

313 part : dtype_u

314 The current particles data.

315 t : float

316 The current time.

318 Returns

319 -------

320 rhs : acceleration

321 Correct right-hand side of type acceleration.

322 """

324 if not isinstance(part, particles):

325 raise ProblemError('something is wrong during build_f, got %s' % type(part))

327 N = self.nparts

329 rhs = acceleration(self.init)

330 for n in range(N):

331 rhs[:, n] = part.q[n] / part.m[n] * (f.elec[:, n] + np.cross(part.vel[:, n], f.magn[:, n]))

333 return rhs

335 # noinspection PyTypeChecker

336 def boris_solver(self, c, dt, old_fields, new_fields, old_parts):

337 r"""

338 The actual Boris solver for static (!) B fields, extended by the c-term.

340 Parameters

341 ----------

342 c : dtype_u

343 The c term gathering the known values from the previous iteration.

344 dt : float

345 The (probably scaled) time step size.

346 old_fields : dtype_f

347 The field values at the previous node :math:m.

348 new_fields : dtype_f

349 The field values at the current node :math:m+1.

350 old_parts : dtype_u

351 The particles at the previous node :math:m.

353 Returns

354 -------

355 vel : particles

356 The velocities at the :math:(m+1)-th node.

357 """

359 N = self.nparts

360 vel = particles.velocity(self.init)

361 self.work_counters['Boris_solver']()

362 Emean = 0.5 * (old_fields.elec + new_fields.elec)

363 for n in range(N):

364 a = old_parts.q[n] / old_parts.m[n]

366 c[:, n] += dt / 2 * a * np.cross(old_parts.vel[:, n], old_fields.magn[:, n] - new_fields.magn[:, n])

368 # pre-velocity, separated by the electric forces (and the c term)

369 vm = old_parts.vel[:, n] + dt / 2 * a * Emean[:, n] + c[:, n] / 2

370 # rotation

371 t = dt / 2 * a * new_fields.magn[:, n]

372 s = 2 * t / (1 + np.linalg.norm(t, 2) ** 2)

373 vp = vm + np.cross(vm + np.cross(vm, t), s)

374 # post-velocity

375 vel[:, n] = vp + dt / 2 * a * Emean[:, n] + c[:, n] / 2

377 return vel