Coverage for pySDC/implementations/hooks/plotting.py: 100%

1from pySDC.core.hooks import Hooks 

2import matplotlib.pyplot as plt 

3 

4 

5class PlottingHook(Hooks): # pragma: no cover 

6 """ 

7 Base class for hooks that plot the solution with the problem's `plot` method, optionally saving each figure. 

8 """ 

9 

10 save_plot = None # Supply a string to the path where you want to save 

11 live_plot = 1e-9 # Supply `None` if you don't want live plotting 

12 

13 def __init__(self): 

14 """Start counting the plots from zero.""" 

15 super().__init__() 

16 self.plot_counter = 0 

17 

18 def pre_run(self, step, level_number): 

19 """ 

20 Get the figure to plot into from the problem's ``get_fig`` method. 

21 

22 Args: 

23 step (pySDC.Step.step): The current step 

24 level_number (int): Number of current level 

25 

26 Returns: 

27 None 

28 """ 

29 super().pre_run(step, level_number) 

30 prob = step.levels[level_number].prob 

31 self.fig = prob.get_fig() 

32 

33 def plot(self, step, level_number, plot_ic=False): 

34 """ 

35 Plot the initial conditions or the solution at the end of the step with the problem's ``plot`` method. The 

36 figure is saved to ``<save_plot>_<counter>.png`` if ``save_plot`` is set, and shown for ``live_plot`` seconds if 

37 that is not ``None``. 

38 

39 Args: 

40 step (pySDC.Step.step): The current step 

41 level_number (int): Number of current level 

42 plot_ic (bool): plot the initial conditions ``u[0]`` instead of the end point 

43 

44 Returns: 

45 None 

46 """ 

47 level = step.levels[level_number] 

48 prob = level.prob 

49 

50 if plot_ic: 

51 u = level.u[0] 

52 t = level.time 

53 else: 

54 level.sweep.compute_end_point() 

55 u = level.uend 

56 t = level.time + level.dt 

57 

58 prob.plot(u=u, t=t, fig=self.fig) 

59 

60 if self.save_plot is not None: 

61 path = f'{self.save_plot}_{self.plot_counter:04d}.png' 

62 self.fig.savefig(path, dpi=100) 

63 self.logger.log(25, f'Saved figure {path!r}.') 

64 

65 if self.live_plot is not None: 

66 plt.pause(self.live_plot) 

67 

68 self.plot_counter += 1 

69 

70 

71class PlotPostStep(PlottingHook): # pragma: no cover 

72 """ 

73 Call a plotting function of the problem after every step 

74 """ 

75 

76 plot_every = 1 

77 

78 def __init__(self): 

79 """Start counting the steps since the last plot from zero.""" 

80 super().__init__() 

81 self.skip_counter = 0 

82 

83 def pre_run(self, step, level_number): 

84 """ 

85 Set up the figure and plot the initial conditions, on the finest level only. 

86 

87 Args: 

88 step (pySDC.Step.step): The current step 

89 level_number (int): Number of current level 

90 

91 Returns: 

92 None 

93 """ 

94 if level_number > 0: 

95 # no figure on coarse levels, but the rest of the chain 

96 return super(PlottingHook, self).pre_run(step, level_number) 

97 super().pre_run(step, level_number) 

98 self.plot(step, level_number, plot_ic=True) 

99 

100 def post_step(self, step, level_number): 

101 """ 

102 Call the plotting function after the step 

103 

104 Args: 

105 step (pySDC.Step.step): The current step 

106 level_number (int): Number of current level 

107 

108 Returns: 

109 None 

110 """ 

111 super().post_step(step, level_number) 

112 if level_number > 0: 

113 return 

114 

115 self.skip_counter += 1 

116 

117 if self.skip_counter % self.plot_every >= 1: 

118 return 

119 

120 self.plot(step, level_number)