Coverage for pySDC/tutorial/step_3/HookClass_Particles.py: 100%

1import numpy as np 

2 

3from pySDC.core.hooks import Hooks 

4 

5 

6class particle_hook(Hooks): 

7 def __init__(self): 

8 """ 

9 Initialization of particles output 

10 """ 

11 super(particle_hook, self).__init__() 

12 

13 def pre_run(self, step, level_number): 

14 """ 

15 Overwrite default routine called before time-loop starts 

16 Args: 

17 step: the current step 

18 level_number: the current level number 

19 """ 

20 super(particle_hook, self).pre_run(step, level_number) 

21 

22 # some abbreviations 

23 L = step.levels[level_number] 

24 part = L.u[0] 

25 N = L.prob.nparts 

26 w = np.array([1, 1, -2]) 

27 

28 # compute (slowly..) the potential at u0 

29 fpot = np.zeros(N) 

30 for i in range(N): 

31 # inner loop, omit ith particle 

32 for j in range(0, i): 

33 dist2 = np.linalg.norm(part.pos[:, i] - part.pos[:, j], 2) ** 2 + L.prob.sig**2 

34 fpot[i] += part.q[j] / np.sqrt(dist2) 

35 for j in range(i + 1, N): 

36 dist2 = np.linalg.norm(part.pos[:, i] - part.pos[:, j], 2) ** 2 + L.prob.sig**2 

37 fpot[i] += part.q[j] / np.sqrt(dist2) 

38 fpot[i] -= L.prob.omega_E**2 * part.m[i] / part.q[i] / 2.0 * np.dot(w, part.pos[:, i] * part.pos[:, i]) 

39 

40 # add up kinetic and potntial contributions to total energy 

41 epot = 0 

42 ekin = 0 

43 for n in range(N): 

44 epot += part.q[n] * fpot[n] 

45 ekin += part.m[n] / 2.0 * np.dot(part.vel[:, n], part.vel[:, n]) 

46 

47 self.add_to_stats( 

48 process=step.status.slot, 

49 time=L.time, 

50 level=L.level_index, 

51 iter=0, 

52 sweep=L.status.sweep, 

53 type='etot', 

54 value=epot + ekin, 

55 ) 

56 

57 def post_step(self, step, level_number): 

58 """ 

59 Default routine called after each iteration 

60 Args: 

61 step: the current step 

62 level_number: the current level number 

63 """ 

64 

65 super(particle_hook, self).post_step(step, level_number) 

66 

67 # some abbreviations 

68 L = step.levels[level_number] 

69 L.sweep.compute_end_point() 

70 part = L.uend 

71 N = L.prob.nparts 

72 w = np.array([1, 1, -2]) 

73 

74 # compute (slowly..) the potential at uend 

75 fpot = np.zeros(N) 

76 for i in range(N): 

77 # inner loop, omit ith particle 

78 for j in range(0, i): 

79 dist2 = np.linalg.norm(part.pos[:, i] - part.pos[:, j], 2) ** 2 + L.prob.sig**2 

80 fpot[i] += part.q[j] / np.sqrt(dist2) 

81 for j in range(i + 1, N): 

82 dist2 = np.linalg.norm(part.pos[:, i] - part.pos[:, j], 2) ** 2 + L.prob.sig**2 

83 fpot[i] += part.q[j] / np.sqrt(dist2) 

84 fpot[i] -= L.prob.omega_E**2 * part.m[i] / part.q[i] / 2.0 * np.dot(w, part.pos[:, i] * part.pos[:, i]) 

85 

86 # add up kinetic and potntial contributions to total energy 

87 epot = 0 

88 ekin = 0 

89 for n in range(N): 

90 epot += part.q[n] * fpot[n] 

91 ekin += part.m[n] / 2.0 * np.dot(part.vel[:, n], part.vel[:, n]) 

92 

93 self.add_to_stats( 

94 process=step.status.slot, 

95 time=L.time, 

96 level=L.level_index, 

97 iter=step.status.iter, 

98 sweep=L.status.sweep, 

99 type='etot', 

100 value=epot + ekin, 

101 ) 

102 

103 return None