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

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 function ``format_index`` class attribute. This should 

85 accept one index and return one string. 

86 

87 You can also give a custom ``logging_condition`` function, 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 function 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 counter = 0 

103 

104 def logging_condition(L): 

105 return True 

106 

107 def process_solution(L): 

108 return {'t': L.time + L.dt, 'u': L.uend.view(np.ndarray)} 

109 

110 def format_index(index): 

111 return f'{index:06d}' 

112 

113 def __init__(self): 

114 super().__init__() 

115 

116 if self.path is None: 

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

118 

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

120 raise ValueError( 

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

122 ) 

123 

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

125 os.mkdir(self.path) 

126 

127 def log_to_file(self, step, level_number, condition, process_solution=None): 

128 if level_number > 0: 

129 return None 

130 

131 L = step.levels[level_number] 

132 

133 if condition: 

134 path = self.get_path(self.counter) 

135 

136 if process_solution: 

137 data = process_solution(L) 

138 else: 

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

140 

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

142 pickle.dump(data, file) 

143 self.logger.info(f'Stored file {path!r}') 

144 

145 type(self).counter += 1 

146 

147 def post_step(self, step, level_number): 

148 L = step.levels[level_number] 

149 self.log_to_file(step, level_number, type(self).logging_condition(L)) 

150 

151 def pre_run(self, step, level_number): 

152 L = step.levels[level_number] 

153 L.uend = L.u[0] 

154 

155 def process_solution(L): 

156 return { 

157 **type(self).process_solution(L), 

158 't': L.time, 

159 } 

160 

161 self.log_to_file(step, level_number, True, process_solution=process_solution) 

162 

163 @classmethod 

164 def get_path(cls, index): 

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

166 

167 @classmethod 

168 def load(cls, index): 

169 path = cls.get_path(index) 

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

171 return pickle.load(file) 

172 

173 

174class LogToFileAfterXs(LogToFile): 

175 r''' 

176 Log to file after certain amount of time has passed instead of after every step 

177 ''' 

178 

179 time_increment = 0 

180 t_next_log = 0 

181 

182 def post_step(self, step, level_number): 

183 L = step.levels[level_number] 

184 

185 if self.t_next_log == 0: 

186 self.t_next_log = self.time_increment 

187 

188 if L.time + L.dt >= self.t_next_log and not step.status.restart: 

189 super().post_step(step, level_number) 

190 self.t_next_log = max([L.time + L.dt, self.t_next_log]) + self.time_increment