Coverage for pySDC/helpers/visualization_tools.py: 100%

1import matplotlib.pyplot as plt 

2import numpy as np 

3 

4from pySDC.helpers.stats_helper import filter_stats 

5 

6 

7# noinspection PyShadowingBuiltins 

8def show_residual_across_simulation(stats, fname='residuals.png'): 

9 """ 

10 Helper routine to visualize the residuals across the simulation (one block of PFASST) 

11 

12 Args: 

13 stats (dict): statistics object from a PFASST run 

14 fname (str): filename 

15 

16 Returns: 

17 matplotlib.figure.Figure: the figure, which is also saved to `fname` 

18 """ 

19 

20 # get residuals of the run 

21 extract_stats = filter_stats(stats, type='residual_post_iteration') 

22 

23 # find boundaries for x-,y- and c-axis as well as arrays 

24 maxprocs = 0 

25 maxiter = 0 

26 minres = 0 

27 maxres = -99 

28 for k, v in extract_stats.items(): 

29 maxprocs = max(maxprocs, k.process) 

30 maxiter = max(maxiter, k.iter) 

31 minres = min(minres, np.log10(v)) 

32 maxres = max(maxres, np.log10(v)) 

33 

34 # grep residuals and put into array 

35 residual = np.zeros((maxiter, maxprocs + 1)) 

36 residual[:] = -99 

37 for k, v in extract_stats.items(): 

38 step = k.process 

39 iter = k.iter 

40 if iter != -1: 

41 residual[iter - 1, step] = np.log10(v) 

42 

43 # large fonts for this figure only, not for every figure made after it 

44 with plt.rc_context( 

45 {'font.size': 30, 'legend.fontsize': 'small', 'xtick.labelsize': 'small', 'ytick.labelsize': 'small'} 

46 ): 

47 fig, ax = plt.subplots(figsize=(15, 10)) 

48 

49 mesh = ax.pcolor(residual.T, cmap=plt.get_cmap('Reds'), vmin=minres, vmax=maxres) 

50 fig.colorbar(mesh, ax=ax).set_label('log10(residual)') 

51 

52 ax.set_xlabel('iteration') 

53 ax.set_ylabel('process') 

54 

55 ax.set_xticks(np.arange(maxiter) + 0.5, minor=False) 

56 ax.set_yticks(np.arange(maxprocs + 1) + 0.5, minor=False) 

57 ax.set_xticklabels(np.arange(maxiter) + 1, minor=False) 

58 ax.set_yticklabels(np.arange(maxprocs + 1), minor=False) 

59 

60 fig.savefig(fname, transparent=True, bbox_inches='tight') 

61 return fig