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

1from pySDC.core.hooks import Hooks 

2import pickle 

3import os 

4import numpy as np 

5from pySDC.helpers.fieldsIO import FieldsIO 

6from pySDC.core.errors import DataError 

7 

8 

9class LogSolution(Hooks): 

10 """ 

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

12 """ 

13 

14 def post_step(self, step, level_number): 

15 """ 

16 Record solution at the end of the step 

17 

18 Args: 

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

20 level_number (int): the current level number 

21 

22 Returns: 

23 None 

24 """ 

25 super().post_step(step, level_number) 

26 

27 L = step.levels[level_number] 

28 L.sweep.compute_end_point() 

29 

30 self.add_to_stats( 

31 process=step.status.slot, 

32 time=L.time + L.dt, 

33 level=L.level_index, 

34 iter=step.status.iter, 

35 sweep=L.status.sweep, 

36 type='u', 

37 value=L.uend, 

38 ) 

39 

40 

41class LogSolutionAfterIteration(Hooks): 

42 """ 

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

44 """ 

45 

46 def post_iteration(self, step, level_number): 

47 """ 

48 Record solution at the end of the iteration 

49 

50 Args: 

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

52 level_number (int): the current level number 

53 

54 Returns: 

55 None 

56 """ 

57 super().post_iteration(step, level_number) 

58 

59 L = step.levels[level_number] 

60 L.sweep.compute_end_point() 

61 

62 self.add_to_stats( 

63 process=step.status.slot, 

64 time=L.time + L.dt, 

65 level=L.level_index, 

66 iter=step.status.iter, 

67 sweep=L.status.sweep, 

68 type='u', 

69 value=L.uend, 

70 ) 

71 

72 

73class LogToPickleFile(Hooks): 

74 r""" 

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

76 

77 Please configure the hook to your liking by setting class attributes on a subclass, so that other runs in the same 

78 process are not affected. You must set a custom path to a directory like so: 

79 

80 ``` 

81 class MyLogToPickleFile(LogToPickleFile): 

82 path = '/my/directory/' 

83 ``` 

84 

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

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

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

88 accept one index and return one string. 

89 

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

91 

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

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

94 

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

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

97 beforehand. 

98 

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

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

101 """ 

102 

103 path = None 

104 file_name = 'solution' 

105 counter = 0 

106 

107 def logging_condition(L): 

108 return True 

109 

110 def process_solution(L): 

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

112 

113 def format_index(index): 

114 return f'{index:06d}' 

115 

116 def __init__(self): 

117 super().__init__() 

118 

119 if self.path is None: 

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

121 

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

123 raise ValueError( 

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

125 ) 

126 

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

128 os.makedirs(self.path, exist_ok=True) 

129 

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

131 if level_number > 0: 

132 return None 

133 

134 L = step.levels[level_number] 

135 

136 if condition: 

137 path = self.get_path(self.counter) 

138 

139 if process_solution: 

140 data = process_solution(L) 

141 else: 

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

143 

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

145 pickle.dump(data, file) 

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

147 

148 type(self).counter += 1 

149 

150 def post_step(self, step, level_number): 

151 L = step.levels[level_number] 

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

153 

154 def pre_run(self, step, level_number): 

155 L = step.levels[level_number] 

156 L.uend = L.u[0] 

157 

158 def process_solution(L): 

159 return { 

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

161 't': L.time, 

162 } 

163 

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

165 

166 @classmethod 

167 def get_path(cls, index): 

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

169 

170 @classmethod 

171 def load(cls, index): 

172 path = cls.get_path(index) 

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

174 return pickle.load(file) 

175 

176 

177class LogToPickleFileAfterXS(LogToPickleFile): 

178 r''' 

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

180 ''' 

181 

182 time_increment = 0 

183 t_next_log = 0 

184 

185 def post_step(self, step, level_number): 

186 L = step.levels[level_number] 

187 

188 if self.t_next_log == 0: 

189 self.t_next_log = self.time_increment 

190 

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

192 super().post_step(step, level_number) 

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

194 

195 def pre_run(self, step, level_number): 

196 L = step.levels[level_number] 

197 L.uend = L.u[0] 

198 

199 def process_solution(L): 

200 return { 

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

202 't': L.time, 

203 } 

204 

205 self.log_to_file(step, level_number, type(self).logging_condition(L), process_solution=process_solution) 

206 

207 

208class LogToFile(Hooks): 

209 filename = 'myRun.pySDC' 

210 time_increment = 0 

211 allow_overwriting = False 

212 counter = 0 # number of stored time points in the file 

213 

214 def __init__(self): 

215 super().__init__() 

216 self.outfile = None 

217 self.t_next_log = 0 

218 FieldsIO.ALLOW_OVERWRITE = self.allow_overwriting 

219 

220 def pre_run(self, step, level_number): 

221 if level_number > 0: 

222 return None 

223 L = step.levels[level_number] 

224 

225 # setup outfile 

226 if os.path.isfile(self.filename) and L.time > 0: 

227 L.prob.setUpFieldsIO() 

228 self.outfile = FieldsIO.fromFile(self.filename) 

229 self.counter = len(self.outfile.times) 

230 self.logger.info( 

231 f'Set up file {self.filename!r} for writing output. This file already contains {self.counter} solutions up to t={self.outfile.times[-1]:.4f}.' 

232 ) 

233 else: 

234 self.outfile = L.prob.getOutputFile(self.filename) 

235 self.logger.info(f'Set up file {self.filename!r} for writing output.') 

236 

237 # write initial conditions 

238 if L.time not in self.outfile.times: 

239 self.outfile.addField(time=L.time, field=L.prob.processSolutionForOutput(L.u[0])) 

240 self.logger.info(f'Written initial conditions at t={L.time:4f} to file') 

241 

242 type(self).counter = len(self.outfile.times) 

243 self.logger.info(f'Will write to disk every {self.time_increment:.4e} time units') 

244 

245 def post_step(self, step, level_number): 

246 if level_number > 0: 

247 return None 

248 

249 L = step.levels[level_number] 

250 

251 if self.t_next_log == 0: 

252 self.t_next_log = L.time + self.time_increment 

253 

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

255 value_exists = True in [abs(me - (L.time + L.dt)) < np.finfo(float).eps * 1000 for me in self.outfile.times] 

256 if value_exists and not self.allow_overwriting: 

257 raise DataError(f'Already have recorded data for time {L.time + L.dt} in this file!') 

258 self.outfile.addField(time=L.time + L.dt, field=L.prob.processSolutionForOutput(L.uend)) 

259 self.logger.info(f'Written solution at t={L.time+L.dt:.4f} to file') 

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

261 type(self).counter = len(self.outfile.times) 

262 

263 def post_run(self, step, level_number): 

264 if level_number > 0: 

265 return None 

266 

267 L = step.levels[level_number] 

268 

269 value_exists = True in [abs(me - (L.time + L.dt)) < np.finfo(float).eps * 1000 for me in self.outfile.times] 

270 if not value_exists: 

271 self.outfile.addField(time=L.time + L.dt, field=L.prob.processSolutionForOutput(L.uend)) 

272 self.logger.info(f'Written solution at t={L.time+L.dt:.4f} to file') 

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

274 type(self).counter = len(self.outfile.times) 

275 

276 @classmethod 

277 def load(cls, index): 

278 data = {} 

279 file = FieldsIO.fromFile(cls.filename) 

280 file_entry = file.readField(idx=index) 

281 data['u'] = file_entry[1] 

282 data['t'] = file_entry[0] 

283 return data