Coverage for pySDC/implementations/problem_classes/PenningTrap_3D.py: 92%
120 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-20 14:51 +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')))
90 self._makeAttributeAndRegister('nparts', localVars=locals(), readOnly=True)
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