Coverage for pySDC/implementations/hooks/log_solution.py: 89%

1from pySDC.core.Hooks import hooks 

2import pickle 

3import os 

4import numpy as np 

5 

6 

7class LogSolution(hooks): 

8 """ 

9 Store the solution at the end of each step as "u". 

10 """ 

11 

12 def post_step(self, step, level_number): 

13 """ 

14 Record solution at the end of the step 

15 

16 Args: 

17 step (pySDC.Step.step): the current step 

18 level_number (int): the current level number 

19 

20 Returns: 

21 None 

22 """ 

23 super().post_step(step, level_number) 

24 

25 L = step.levels[level_number] 

26 L.sweep.compute_end_point() 

27 

28 self.add_to_stats( 

29 process=step.status.slot, 

30 time=L.time + L.dt, 

31 level=L.level_index, 

32 iter=step.status.iter, 

33 sweep=L.status.sweep, 

34 type='u', 

35 value=L.uend, 

36 ) 

37 

38 

39class LogSolutionAfterIteration(hooks): 

40 """ 

41 Store the solution at the end of each iteration as "u". 

42 """ 

43 

44 def post_iteration(self, step, level_number): 

45 """ 

46 Record solution at the end of the iteration 

47 

48 Args: 

49 step (pySDC.Step.step): the current step 

50 level_number (int): the current level number 

51 

52 Returns: 

53 None 

54 """ 

55 super().post_iteration(step, level_number) 

56 

57 L = step.levels[level_number] 

58 L.sweep.compute_end_point() 

59 

60 self.add_to_stats( 

61 process=step.status.slot, 

62 time=L.time + L.dt, 

63 level=L.level_index, 

64 iter=step.status.iter, 

65 sweep=L.status.sweep, 

66 type='u', 

67 value=L.uend, 

68 ) 

69 

70 

71class LogToFile(hooks): 

72 r""" 

73 Hook for logging the solution to file after the step using pickle. 

74 

75 Please configure the hook to your liking by manipulating class attributes. 

76 You must set a custom path to a directory like so: 

77 

78 ``` 

79 LogToFile.path = '/my/directory/' 

80 ``` 

81 

82 Keep in mind that the hook will overwrite files without warning! 

83 You can give a custom file name by setting the ``file_name`` class attribute and give a custom way of rendering the 

84 index associated with individual files by giving a different lambda function ``format_index`` class attribute. This 

85 lambda should accept one index and return one string. 

86 

87 You can also give a custom ``logging_condition`` lambda, accepting the current level if you want to log selectively. 

88 

89 Importantly, you may need to change ``process_solution``. By default, this will return a numpy view of the solution. 

90 Of course, if you are not using numpy, you need to change this. Again, this is a lambda accepting the level. 

91 

92 After the fact, you can use the classmethod `get_path` to get the path to a certain data or the `load` function to 

93 directly load the solution at a given index. Just configure the hook like you did when you recorded the data 

94 beforehand. 

95 

96 Finally, be aware that using this hook with MPI parallel runs may lead to different tasks overwriting files. Make 

97 sure to give a different `file_name` for each task that writes files. 

98 """ 

99 

100 path = None 

101 file_name = 'solution' 

102 logging_condition = lambda L: True 

103 process_solution = lambda L: {'t': L.time + L.dt, 'u': L.uend.view(np.ndarray)} 

104 format_index = lambda index: f'{index:06d}' 

105 

106 def __init__(self): 

107 super().__init__() 

108 self.counter = 0 

109 

110 if self.path is None: 

111 raise ValueError('Please set a path for logging as the class attribute `LogToFile.path`!') 

112 

113 if os.path.isfile(self.path): 

114 raise ValueError( 

115 f'{self.path!r} is not a valid path to log to because a file of the same name exists. Please supply a directory' 

116 ) 

117 

118 if not os.path.isdir(self.path): 

119 os.mkdir(self.path) 

120 

121 def post_step(self, step, level_number): 

122 if level_number > 0: 

123 return None 

124 

125 L = step.levels[level_number] 

126 

127 if type(self).logging_condition(L): 

128 path = self.get_path(self.counter) 

129 data = type(self).process_solution(L) 

130 

131 with open(path, 'wb') as file: 

132 pickle.dump(data, file) 

133 

134 self.counter += 1 

135 

136 @classmethod 

137 def get_path(cls, index): 

138 return f'{cls.path}/{cls.file_name}_{cls.format_index(index)}.pickle' 

139 

140 @classmethod 

141 def load(cls, index): 

142 path = cls.get_path(index) 

143 with open(path, 'rb') as file: 

144 return pickle.load(file)