Coverage for pySDC/implementations/problem_classes/PenningTrap_3D.py: 99%
108 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 13:41 +0000
1import numpy as np
3from pySDC.core.errors import ProblemError
4from pySDC.core.problem import Problem, WorkCounter
5from pySDC.implementations.datatype_classes.particles import particles, fields, acceleration
8# noinspection PyUnusedLocal
9class penningtrap(Problem):
10 r"""
11 Charged particles in a 3D Penning trap with Coulomb interaction, a second-order problem for the Boris integrator.
13 This class implements a standard Penning trap problem on the time interval :math:`[0, t_{end}]`
14 fully investigated in [1]_. The equations are given by the following equation of motion
16 .. math::
17 \frac{dv}{dt}=f(x,v)=\alpha[E(x,t)+v\times B(x,t)],
19 .. math::
20 \frac{dx}{dt}=v
22 with the particles :math:`x, v\in \mathbb{R}^{3}`. For the penning trap problem, the other parameters are given by
23 the constant magnetic field :math:`B=\frac{\omega_{B}}{\alpha}\cdot \hat{e_{z}}\in \mathbb{R}^{3}`
24 along the :math:`z`-axis with the particle's charge-to-mass ratio :math:`\alpha=\frac{q}{m}` so that
26 .. math::
27 v\times B=\frac{\omega_{B}}{\alpha}\left(
28 \begin{matrix}
29 0 & 1 & 0\\
30 -1 & 0 & 0\\
31 0 & 0 & 0
32 \end{matrix}
33 \right)v.
35 The electric field :math:`E(x_{i})=E_{ext}(x_{i})+E_{int}(x_{i})\in \mathbb{R}^{3}` where
37 .. math::
38 E_{ext}(x_{i})=-\epsilon\frac{\omega_{E}^{2}}{\alpha}\left(
39 \begin{matrix}
40 1 & 0 & 0\\
41 0 & 1 & 0\\
42 0 & 0 & -2
43 \end{matrix}
44 \right)x_{i}
46 and the inter-particle Coulomb interaction
48 .. math::
49 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}}
51 with the smoothing parameter :math:`\lambda>0`.
52 The exact solution also given for the single particle penning trap more detailed [1]_, [2]_.
53 For to solve nonlinear equation of system, Boris trick is used (see [2]_).
55 Parameters
56 ----------
57 omega_B : float
58 Amplitude of magnetic field.
59 omega_E : float
60 Amplitude of electric field.
61 u0 : np.1darray
62 Initial condition ``[pos, vel, [q], [m]]``: position and velocity (three components each) around which the
63 particles are placed, and the particles' charge and mass (so far only :math:`q = m = 1` is implemented).
64 nparts : int
65 The number of particles.
66 sig : float
67 The smoothing parameter :math:`\lambda>0`.
69 Attributes
70 ----------
71 work_counter : dict
72 Counts the calls of the right-hand side, and calls of the Boris solver.
74 References
75 ----------
76 .. [1] F. Penning. Die Glimmentladung bei niedrigem Druck zwischen koaxialen Zylindern in einem axialen Magnetfeld.
77 Physica. Vol. 3 (1936).
78 .. [2] Mathias Winkel, Robert Speck and Daniel Ruprecht. A high-order Boris integrator.
79 Journal of Computational Physics (2015).
80 """
82 dtype_u = particles
83 dtype_f = fields
85 def __init__(self, omega_B, omega_E, u0, nparts, sig):
86 """Initialization routine"""
87 # invoke super init, passing nparts, dtype_u and dtype_f
88 super().__init__(((3, nparts), None, np.dtype('float64')))
89 self._makeAttributeAndRegister('nparts', localVars=locals(), readOnly=True)
90 self._makeAttributeAndRegister('omega_B', 'omega_E', 'u0', 'sig', localVars=locals())
91 self.work_counters['Boris_solver'] = WorkCounter()
92 self.work_counters['rhs'] = WorkCounter()
94 @staticmethod
95 def fast_interactions(N, pos, sig, q):
96 r"""
97 Computes the fast interactions.
99 All pairs at once, with NumPy broadcasting: three times faster than the numba-compiled loop this
100 replaces, and with no compiler to install. The pairwise differences take :math:`3N^2` values, a few
101 megabytes for the hundreds of particles a direct sum is meant for. The self-interaction vanishes
102 since :math:`\lambda > 0`.
104 Parameters
105 ----------
106 N : int
107 Number of particles.
108 pos : np.2darray
109 Position.
110 sig : float
111 The smoothing parameter :math:`\lambda > 0`.
112 q : np.1darray
113 Particle's charge.
115 Returns
116 -------
117 Efield : np.2darray
118 The internal E field for each particle.
119 """
120 pos = np.asarray(pos) # a plain array: the mesh wrapper would cost more than this arithmetic
121 diff = pos[:, :, None] - pos[:, None, :] # (3, N, N): position of particle i minus that of particle j
122 dist2 = diff[0] ** 2 + diff[1] ** 2 + diff[2] ** 2 + sig**2
123 return (q * diff / dist2**1.5).sum(axis=2)
125 def get_interactions(self, part):
126 r"""
127 Routine to compute the particle-particle interaction, assuming :math:`q = 1` for all particles.
129 Parameters
130 ----------
131 part : dtype_u
132 The particles.
134 Returns
135 -------
136 Efield : np.ndarray
137 The internal E field for each particle.
138 """
140 N = self.nparts
142 Efield = self.fast_interactions(N, part.pos, self.sig, part.q)
144 return Efield
146 def eval_f(self, part, t):
147 """
148 Routine to compute the E and B fields (named f for consistency with the original PEPC version).
150 Parameters
151 ----------
152 part : dtype_u
153 The particles.
154 t : float
155 Current time of the particles (not used here).
157 Returns
158 -------
159 f : dtype_f
160 Fields for the particles (internal and external), i.e., the right-hand side of the problem.
161 """
163 self.work_counters['rhs']()
164 try:
165 _unused = penningtrap.Harmonic_oscillator
166 Emat = np.diag([0, 0, -1])
167 except AttributeError:
168 Emat = np.diag([1, 1, -2])
170 f = self.dtype_f(self.init)
172 f.elec[:] = self.get_interactions(part) + self.omega_E**2 / (part.q / part.m) * (Emat @ part.pos)
173 f.magn[:] = self.omega_B * np.array([0, 0, 1])[:, None]
175 return f
177 # TODO : Warning, this should be moved to u_exact(t=0) !
178 def u_init(self):
179 """
180 Routine to compute the starting values for the particles.
182 The first particle starts at ``u0``. The others are shifted from it by random amounts in :math:`[-1, 0)` in
183 each position component and in :math:`[-5, -4)` in each velocity component, drawn after seeding
184 ``np.random`` with the number of particles.
186 Returns
187 -------
188 u : dtype_u
189 Particle set filled with initial data.
190 """
192 u0 = self.u0
193 N = self.nparts
195 u = self.dtype_u(self.init)
197 if u0[2][0] != 1 or u0[3][0] != 1:
198 raise ProblemError('so far only q = m = 1 is implemented')
200 # set first particle to u0
201 u.pos[0, 0] = u0[0][0]
202 u.pos[1, 0] = u0[0][1]
203 u.pos[2, 0] = u0[0][2]
204 u.vel[0, 0] = u0[1][0]
205 u.vel[1, 0] = u0[1][1]
206 u.vel[2, 0] = u0[1][2]
208 u.q[0] = u0[2][0]
209 u.m[0] = u0[3][0]
211 # initialize random seed
212 np.random.seed(N)
214 comx = u.pos[0, 0]
215 comy = u.pos[1, 0]
216 comz = u.pos[2, 0]
218 for n in range(1, N):
219 # draw 3 random variables in [-1, 0) to shift positions
220 r = np.random.random_sample(3) - 1
221 u.pos[0, n] = r[0] + u0[0][0]
222 u.pos[1, n] = r[1] + u0[0][1]
223 u.pos[2, n] = r[2] + u0[0][2]
225 # draw 3 random variables in [-5, -4) to shift velocities
226 r = np.random.random_sample(3) - 5
227 u.vel[0, n] = r[0] + u0[1][0]
228 u.vel[1, n] = r[1] + u0[1][1]
229 u.vel[2, n] = r[2] + u0[1][2]
231 u.q[n] = u0[2][0]
232 u.m[n] = u0[3][0]
234 # gather positions to check center
235 comx += u.pos[0, n]
236 comy += u.pos[1, n]
237 comz += u.pos[2, n]
239 # print('Center of positions:',comx/N,comy/N,comz/N)
241 return u
243 def u_exact(self, t):
244 r"""
245 Routine to compute the exact trajectory at time :math:`t` (only for single-particle setup).
247 Parameters
248 ----------
249 t : float
250 Current time of the exact trajectory.
252 Returns
253 -------
254 u : dtype_u
255 Particle type containing the exact position and velocity.
256 """
258 # some abbreviations
259 wE = self.omega_E
260 wB = self.omega_B
261 N = self.nparts
262 u0 = self.u0
264 if N != 1:
265 raise ProblemError('u_exact is only valid for a single particle')
267 u = self.dtype_u(((3, 1), self.init[1], self.init[2]))
269 wbar = np.sqrt(2) * wE
271 # position and velocity in z direction is easy to compute
272 u.pos[2, 0] = u0[0][2] * np.cos(wbar * t) + u0[1][2] / wbar * np.sin(wbar * t)
273 u.vel[2, 0] = -u0[0][2] * wbar * np.sin(wbar * t) + u0[1][2] * np.cos(wbar * t)
275 # define temp. variables to compute complex position
276 Op = 1 / 2 * (wB + np.sqrt(wB**2 - 4 * wE**2))
277 Om = 1 / 2 * (wB - np.sqrt(wB**2 - 4 * wE**2))
278 Rm = (Op * u0[0][0] + u0[1][1]) / (Op - Om)
279 Rp = u0[0][0] - Rm
280 Im = (Op * u0[0][1] - u0[1][0]) / (Op - Om)
281 Ip = u0[0][1] - Im
283 # compute position in complex notation
284 w = (Rp + Ip * 1j) * np.exp(-Op * t * 1j) + (Rm + Im * 1j) * np.exp(-Om * t * 1j)
285 # compute velocity as time derivative of the position
286 dw = -1j * Op * (Rp + Ip * 1j) * np.exp(-Op * t * 1j) - 1j * Om * (Rm + Im * 1j) * np.exp(-Om * t * 1j)
288 # get the appropriate real and imaginary parts
289 u.pos[0, 0] = w.real
290 u.vel[0, 0] = dw.real
291 u.pos[1, 0] = w.imag
292 u.vel[1, 0] = dw.imag
294 return u
296 def build_f(self, f, part, t):
297 """
298 Helper function to assemble the correct right-hand side out of B and E field.
300 Parameters
301 ----------
302 f : dtype_f
303 The field values.
304 part : dtype_u
305 The current particles data.
306 t : float
307 The current time.
309 Returns
310 -------
311 rhs : acceleration
312 Correct right-hand side of type acceleration.
313 """
315 if not isinstance(part, particles):
316 raise ProblemError('something is wrong during build_f, got %s' % type(part))
318 rhs = acceleration(self.init)
319 rhs[:] = part.q / part.m * (f.elec + np.cross(part.vel, f.magn, axis=0))
321 return rhs
323 # noinspection PyTypeChecker
324 def boris_solver(self, c, dt, old_fields, new_fields, old_parts):
325 r"""
326 The actual Boris solver for static (!) B fields, extended by the c-term.
328 Parameters
329 ----------
330 c : dtype_u
331 The c term gathering the known values from the previous iteration.
332 dt : float
333 The (probably scaled) time step size.
334 old_fields : dtype_f
335 The field values at the previous node :math:`m`.
336 new_fields : dtype_f
337 The field values at the current node :math:`m+1`.
338 old_parts : dtype_u
339 The particles at the previous node :math:`m`.
341 Returns
342 -------
343 vel : particles
344 The velocities at the :math:`(m+1)`-th node.
345 """
347 vel = particles.velocity(self.init)
348 self.work_counters['Boris_solver']()
350 # All particles at once: positions and velocities are (3, N), charge and mass (N,). On plain arrays, since
351 # the mesh wrapper around each of these small operations would cost more than the operation itself. The
352 # view of `c` still writes into the caller's array.
353 v_old, B_old, B_new = np.asarray(old_parts.vel), np.asarray(old_fields.magn), np.asarray(new_fields.magn)
354 c_term = np.asarray(c)
355 Emean = 0.5 * (np.asarray(old_fields.elec) + np.asarray(new_fields.elec))
356 a = old_parts.q / old_parts.m
358 c_term += dt / 2 * a * np.cross(v_old, B_old - B_new, axis=0)
360 # pre-velocity, separated by the electric forces (and the c term)
361 vm = v_old + dt / 2 * a * Emean + c_term / 2
362 # rotation
363 t = dt / 2 * a * B_new
364 s = 2 * t / (1 + np.linalg.norm(t, 2, axis=0) ** 2)
365 vp = vm + np.cross(vm + np.cross(vm, t, axis=0), s, axis=0)
366 # post-velocity
367 vel[:] = vp + dt / 2 * a * Emean + c_term / 2
369 return vel