Part B: Adding statistics#
Now we extend the statistics with entries of our own. To make things more interesting (and complicated), this part also brings a new problem, a new sweeper and new data types:
PenningTrap_3D: particles in a Penning trap, held in place by an electric and a magnetic field,boris_2nd_order: SDC for second-order problems with the Boris method, see this paper,particlesfor positions, velocities, charges and masses, andfieldsfor the electric and magnetic field.
An important measure for this kind of problem is the total energy of the system, which we would like to compute after each time step. That is a job for a new hook.
A hook of our own#
Hooks are called by the controller at fixed points of a run: before it, before and after each step, iteration
or sweep, and after it. particle_hook, in HookClass_Particles.py next to this tutorial, computes the total
energy before the run and after each step and adds it to the statistics with add_to_stats, as type 'etot'.
This is what it does after a step:
def post_step(self, step, level_number):
"""
Overwrite default routine called after each step
Args:
step: the current step
level_number: the current level number
"""
super(particle_hook, self).post_step(step, level_number)
# some abbreviations
L = step.levels[level_number]
L.sweep.compute_end_point()
part = L.uend
N = L.prob.nparts
w = np.array([1, 1, -2])
# compute (slowly..) the potential at uend
fpot = np.zeros(N)
for i in range(N):
# inner loop, omit ith particle
for j in range(0, i):
dist2 = np.linalg.norm(part.pos[:, i] - part.pos[:, j], 2) ** 2 + L.prob.sig**2
fpot[i] += part.q[j] / np.sqrt(dist2)
for j in range(i + 1, N):
dist2 = np.linalg.norm(part.pos[:, i] - part.pos[:, j], 2) ** 2 + L.prob.sig**2
fpot[i] += part.q[j] / np.sqrt(dist2)
fpot[i] -= L.prob.omega_E**2 * part.m[i] / part.q[i] / 2.0 * np.dot(w, part.pos[:, i] * part.pos[:, i])
# add up kinetic and potntial contributions to total energy
epot = 0
ekin = 0
for n in range(N):
epot += part.q[n] * fpot[n]
ekin += part.m[n] / 2.0 * np.dot(part.vel[:, n], part.vel[:, n])
self.add_to_stats(
process=step.status.slot,
time=L.time,
level=L.level_index,
iter=step.status.iter,
sweep=L.status.sweep,
type='etot',
value=epot + ekin,
)
return None
Our hook does not replace pySDC’s default statistics: those come from DefaultHooks, which the controller always
adds next to ours. The call to super() first still matters: the base class keeps track of restarted steps, and
add_to_stats writes that into the key of each entry.
import numpy as np
from pathlib import Path
from pySDC.helpers.stats_helper import get_list_of_types, get_sorted
from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI
from pySDC.implementations.problem_classes.PenningTrap_3D import penningtrap
from pySDC.implementations.sweeper_classes.boris_2nd_order import boris_2nd_order
from pySDC.tutorial.step_3.HookClass_Particles import particle_hook
# initialize level parameters
level_params = {'restol': 1e-08, 'dt': 1.0 / 16}
# initialize sweeper parameters
sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': 3}
# initialize problem parameters for the Penning trap
problem_params = {
'omega_E': 4.9, # E-field frequency
'omega_B': 25.0, # B-field frequency
'u0': np.array([[10, 0, 0], [100, 0, 100], [1], [1]], dtype=object), # initial position, velocity, charge, mass
'nparts': 1, # number of particles in the trap
'sig': 0.1, # smoothing parameter for the forces
}
# initialize step parameters
step_params = {'maxiter': 20}
# initialize controller parameters
controller_params = {
'hook_class': particle_hook, # specialized hook class for more statistics and output
'log_to_file': True,
'fname': 'data/step_3_B_out.txt',
}
# Fill description dictionary for easy hierarchy creation
description = {
'problem_class': penningtrap,
'problem_params': problem_params,
'sweeper_class': boris_2nd_order,
'sweeper_params': sweeper_params,
'level_params': level_params,
'step_params': step_params,
}
Path("data").mkdir(parents=True, exist_ok=True)
The hook goes into the controller parameters. As in Step 2, the controller prints its setup when we create it, and the log of the run:
# instantiate the controller
controller = controller_nonMPI(num_procs=1, controller_params=controller_params, description=description)
# set time parameters: a single step
t0 = 0.0
Tend = level_params['dt']
# get initial values on finest level
P = controller.MS[0].levels[0].prob
uinit = P.u_init()
# call main function to get things done...
uend, stats = controller.run(u0=uinit, t0=t0, Tend=Tend)
controller - INFO: Welcome to the one and only, really very astonishing and 87.3% bug free
_____ _____ _____
/ ____| __ \ / ____|
_ __ _ _| (___ | | | | |
| '_ \| | | |\___ \| | | | |
| |_) | |_| |____) | |__| | |____
| .__/ \__, |_____/|_____/ \_____|
| | __/ |
|_| |___/
controller - INFO: Setup overview (--> user-defined, -> dependency) -- BEGIN
controller - INFO: ----------------------------------------------------------------------------------------------------
Controller: <class 'pySDC.implementations.controller_classes.controller_nonMPI.controller_nonMPI'>
all_to_done = False
dump_setup = True
--> fname = data/step_3_B_out.txt
--> hook_class = [<class 'pySDC.core.default_hook.DefaultHooks'>, <class 'pySDC.core.timings.CPUTimings'>, <class 'pySDC.tutorial.step_3.HookClass_Particles.particle_hook'>]
--> log_to_file = True
logger_level = 20
mssdc_jac = True
predict_type = None
use_iteration_estimator = False
Step: <class 'pySDC.core.step.Step'>
--> maxiter = 20
Number of steps: None
Level: <class 'pySDC.core.level.Level'>
Level 0
--> dt = 0.0625
dt_initial = 0.0625
nsweeps = 1
residual_type = full_abs
--> restol = 1e-08
--> Problem: <class 'pySDC.implementations.problem_classes.PenningTrap_3D.penningtrap'>
--> nparts = 1
--> omega_B = 25.0
--> omega_E = 4.9
--> sig = 0.1
--> u0 = [list([10, 0, 0]) list([100, 0, 100]) list([1]) list([1])]
-> Data type u: <class 'pySDC.implementations.datatype_classes.particles.particles'>
-> Data type f: <class 'pySDC.implementations.datatype_classes.particles.fields'>
--> Sweeper: <class 'pySDC.implementations.sweeper_classes.boris_2nd_order.boris_2nd_order'>
QE = EE
QI = IE
do_coll_update = False
initial_guess = spread
--> num_nodes = 3
--> quad_type = RADAU-RIGHT
skip_residual_computation = ()
-> Collocation: <class 'pySDC.core.collocation.CollBase'>
Active convergence controllers:
| # | order | convergence controller
----+----+-------+---------------------------------------------------------------------------------------
| 0 | 95 | BasicRestartingNonMPI
-> | 1 | 100 | SpreadStepSizesBlockwiseNonMPI
| 2 | 200 | CheckConvergence
controller - INFO: ----------------------------------------------------------------------------------------------------
controller - INFO: Setup overview (--> user-defined, -> dependency) -- END
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 1 -- Sweep: 1 -- residual: 3.53203678e+00
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 2 -- Sweep: 1 -- residual: 2.09852117e-01
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 3 -- Sweep: 1 -- residual: 3.50301513e-02
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 4 -- Sweep: 1 -- residual: 4.67724741e-03
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 5 -- Sweep: 1 -- residual: 7.95583202e-04
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 6 -- Sweep: 1 -- residual: 1.11405073e-04
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 7 -- Sweep: 1 -- residual: 1.26902403e-05
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 8 -- Sweep: 1 -- residual: 1.16534542e-06
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 9 -- Sweep: 1 -- residual: 1.66967993e-07
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 10 -- Sweep: 1 -- residual: 2.09408029e-08
hooks - INFO: Process 0 on time 0.000000 at stage IT_FINE: Level: 0 -- Iteration: 11 -- Sweep: 1 -- residual: 2.17123386e-09
hooks - INFO: Finished run after 1.54e+00s
Particles#
The solution is a particles data type. Its parts are separate arrays, one column per particle:
print('position:', uend.pos.T, '\nvelocity:', uend.vel.T, '\ncharge:', uend.q, ' mass:', uend.m)
position: [[14.44301633 -4.24938502 6.05642416]]
velocity: [[ 12.81766217 -113.34810048 90.76670718]]
charge: [1.] mass: [1.]
Our statistics#
Our type 'etot' now shows up among the others, and get_sorted treats it like any other:
print('etot is registered:', 'etot' in get_list_of_types(stats))
# filter statistics type (etot)
energy = get_sorted(stats, type='etot', sortby='iter')
# get base energy and show difference
base_energy = energy[0][1]
for item in energy:
print(
'Total energy and deviation in iteration %2i: %12.10f -- %12.8e'
% (item[0], item[1], abs(base_energy - item[1]))
)
etot is registered: True
Total energy and deviation in iteration 0: 8799.5000000000 -- 0.00000000e+00
Total energy and deviation in iteration 11: 8785.0038936088 -- 1.44961064e+01
Iteration 0 is the energy the hook computed before the run, the other one after the step, at the iteration it converged in. For this single particle the exact solution is known, so we can also check the position:
# compute error compared to know exact solution for one particle
uex = P.u_exact(Tend)
err = np.linalg.norm(uex.pos - uend.pos, np.inf) / np.linalg.norm(uex.pos, np.inf)
print(f'relative error of the position: {err:.3e}')
relative error of the position: 4.290e-04
The position is accurate, but the energy has changed by about 14.5 out of 8800 in a single step. We look into that in Part C.
Important things to note
A custom hook calls
super()in every method it overrides, so that its entries are labelled correctly when a step is restarted, e.g. by adaptivity.User-defined statistics can also come from the problem class: give it an attribute (e.g. the number of GMRES iterations of its spatial solver) and read it in the hook through the level, as
L.prob.
The checks the tests run:
assert abs(base_energy - energy[-1][1]) < 15, f'ERROR: energy deviated too much, got {base_energy - energy[-1][1]}'
assert err < 5e-04, f"ERROR: solution is not as exact as expected, got {err}"