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

## 49 statements

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

1import numpy as np

3from pySDC.core.hooks import Hooks

6class particle_hook(Hooks):

7 def __init__(self):

8 """

9 Initialization of particles output

10 """

11 super(particle_hook, self).__init__()

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)

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])

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])

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])

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 )

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 """

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

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])

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])

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])

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 )

103 return None