Coverage for pySDC/tutorial/step_4/PenningTrap_3D_coarse.py: 100%

1import numpy as np 

2 

3from pySDC.implementations.problem_classes.PenningTrap_3D import penningtrap 

4 

5 

6class penningtrap_coarse(penningtrap): 

7 """ 

8 Coarse level problem description class, will only overwrite what is needed 

9 """ 

10 

11 def eval_f(self, part, t): 

12 """ 

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

14 

15 Args: 

16 t: current time (not used here) 

17 part: the particles 

18 Returns: 

19 Fields for the particles (external only) 

20 """ 

21 

22 N = self.nparts 

23 

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

25 f = self.dtype_f(self.init, val=0.0) 

26 

27 # only compute external forces here: O(N) instead of O(N*N) 

28 for n in range(N): 

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

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

31 

32 return f