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

1# --- 

2# jupyter: 

3# jupytext: 

4# formats: py:percent 

5# kernelspec: 

6# display_name: Python 3 

7# name: python3 

8# --- 

9 

10# %% [markdown] 

11# # Part B: Adding statistics 

12# 

13# Now we extend the statistics with entries of our own. To make things more interesting (and complicated), this 

14# part also brings a new problem, a new sweeper and new data types: 

15# 

16# - `PenningTrap_3D`: particles in a Penning trap, held in place by an electric and a magnetic field, 

17# - `boris_2nd_order`: SDC for second-order problems with the Boris method, see 

18# [this paper](http://dx.doi.org/10.1016/j.jcp.2015.04.022), 

19# - `particles` for positions, velocities, charges and masses, and `fields` for the electric and magnetic field. 

20# 

21# An important measure for this kind of problem is the total energy of the system, which we would like to compute 

22# after each time step. That is a job for a new **hook**. 

23# 

24# ## A hook of our own 

25# 

26# Hooks are called by the controller at fixed points of a run: before it, before and after each step, iteration 

27# or sweep, and after it. `particle_hook`, in `HookClass_Particles.py` next to this tutorial, computes the total 

28# energy before the run and after each step and adds it to the statistics with `add_to_stats`, as type `'etot'`. 

29# This is what it does after a step: 

30# 

31# :::{literalinclude} HookClass_Particles.py 

32# :pyobject: particle_hook.post_step 

33# ::: 

34# 

35# Our hook does not replace pySDC's default statistics: those come from `DefaultHooks`, which the controller always 

36# adds next to ours. The call to `super()` first still matters: the base class keeps track of restarted steps, and 

37# `add_to_stats` writes that into the key of each entry. 

38 

39# %% 

40import numpy as np 

41from pathlib import Path 

42 

43from pySDC.helpers.stats_helper import get_list_of_types, get_sorted 

44from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

45from pySDC.implementations.problem_classes.PenningTrap_3D import penningtrap 

46from pySDC.implementations.sweeper_classes.boris_2nd_order import boris_2nd_order 

47from pySDC.tutorial.step_3.HookClass_Particles import particle_hook 

48 

49# initialize level parameters 

50level_params = {'restol': 1e-08, 'dt': 1.0 / 16} 

51 

52# initialize sweeper parameters 

53sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': 3} 

54 

55# initialize problem parameters for the Penning trap 

56problem_params = { 

57 'omega_E': 4.9, # E-field frequency 

58 'omega_B': 25.0, # B-field frequency 

59 'u0': np.array([[10, 0, 0], [100, 0, 100], [1], [1]], dtype=object), # initial position, velocity, charge, mass 

60 'nparts': 1, # number of particles in the trap 

61 'sig': 0.1, # smoothing parameter for the forces 

62} 

63 

64# initialize step parameters 

65step_params = {'maxiter': 20} 

66 

67# initialize controller parameters 

68controller_params = { 

69 'hook_class': particle_hook, # specialized hook class for more statistics and output 

70 'log_to_file': True, 

71 'fname': 'data/step_3_B_out.txt', 

72} 

73 

74# Fill description dictionary for easy hierarchy creation 

75description = { 

76 'problem_class': penningtrap, 

77 'problem_params': problem_params, 

78 'sweeper_class': boris_2nd_order, 

79 'sweeper_params': sweeper_params, 

80 'level_params': level_params, 

81 'step_params': step_params, 

82} 

83 

84Path("data").mkdir(parents=True, exist_ok=True) 

85 

86# %% [markdown] 

87# The hook goes into the controller parameters. As in [Step 2](../step_2/C_using_pySDCs_frontend), the controller 

88# prints its setup when we create it, and the log of the run: 

89 

90# %% tags=["scroll-output"] 

91# instantiate the controller 

92controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description) 

93 

94# set time parameters: a single step 

95t0 = 0.0 

96Tend = level_params['dt'] 

97 

98# get initial values on finest level 

99P = controller.MS[0].levels[0].prob 

100uinit = P.u_init() 

101 

102# call main function to get things done... 

103uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend) 

104 

105# %% [markdown] 

106# ## Particles 

107# 

108# The solution is a `particles` data type. Its parts are separate arrays, one column per particle: 

109 

110# %% 

111print('position:', uend.pos.T, '\nvelocity:', uend.vel.T, '\ncharge:', uend.q, ' mass:', uend.m) 

112 

113# %% [markdown] 

114# ## Our statistics 

115# 

116# Our type `'etot'` now shows up among the others, and `get_sorted` treats it like any other: 

117 

118# %% 

119print('etot is registered:', 'etot' in get_list_of_types(stats)) 

120 

121# filter statistics type (etot) 

122energy = get_sorted(stats, type='etot', sortby='iter') 

123 

124# get base energy and show difference 

125base_energy = energy[0][1] 

126for item in energy: 

127 print( 

128 'Total energy and deviation in iteration %2i: %12.10f -- %12.8e' 

129 % (item[0], item[1], abs(base_energy - item[1])) 

130 ) 

131 

132# %% [markdown] 

133# Iteration 0 is the energy the hook computed before the run, the other one after the step, at the iteration it 

134# converged in. For this single particle the exact solution is known, so we can also check the position: 

135 

136# %% 

137# compute error compared to know exact solution for one particle 

138uex = P.u_exact(Tend) 

139err = np.linalg.norm(uex.pos - uend.pos, np.inf) / np.linalg.norm(uex.pos, np.inf) 

140print(f'relative error of the position: {err:.3e}') 

141 

142# %% [markdown] 

143# The position is accurate, but the energy has changed by about 14.5 out of 8800 in a single step. We look into 

144# that in [Part C](C_study_collocations). 

145# 

146# :::{admonition} Important things to note 

147# - A custom hook calls `super()` in every method it overrides, so that its entries are labelled correctly when a 

148# step is restarted, e.g. by adaptivity. 

149# - User-defined statistics can also come from the problem class: give it an attribute (e.g. the number of GMRES 

150# iterations of its spatial solver) and read it in the hook through the level, as `L.prob`. 

151# ::: 

152# 

153# The checks the tests run: 

154 

155# %% 

156assert abs(base_energy - energy[-1][1]) < 15, f'ERROR: energy deviated too much, got {base_energy - energy[-1][1]}' 

157assert err < 5e-04, f"ERROR: solution is not as exact as expected, got {err}"