Coverage for pySDC/implementations/hooks/log_solution.py: 97%
144 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
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
9class LogSolution(Hooks):
10 """
11 Store the solution at the end of each step as "u".
12 """
14 def post_step(self, step, level_number):
15 """
16 Record solution at the end of the step
18 Args:
19 step (pySDC.Step.step): the current step
20 level_number (int): the current level number
22 Returns:
23 None
24 """
25 super().post_step(step, level_number)
27 L = step.levels[level_number]
28 L.sweep.compute_end_point()
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 )
41class LogSolutionAfterIteration(Hooks):
42 """
43 Store the solution at the end of each iteration as "u".
44 """
46 def post_iteration(self, step, level_number):
47 """
48 Record solution at the end of the iteration
50 Args:
51 step (pySDC.Step.step): the current step
52 level_number (int): the current level number
54 Returns:
55 None
56 """
57 super().post_iteration(step, level_number)
59 L = step.levels[level_number]
60 L.sweep.compute_end_point()
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 )
73class LogToPickleFile(Hooks):
74 r"""
75 Hook for logging the solution to file after the step using pickle.
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:
80 .. code-block:: python
82 class MyLogToPickleFile(LogToPickleFile):
83 path = '/my/directory/'
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.
90 You can also give a custom ``logging_condition`` function, accepting the current level if you want to log selectively.
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.
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.
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 """
103 path = None
104 file_name = 'solution'
105 counter = 0
107 def logging_condition(L):
108 """
109 Decide whether to log after a step. Called on the class, not an instance; the default logs every step.
111 Args:
112 L (pySDC.Level.level): the level to log
114 Returns:
115 bool: whether to write a file
116 """
117 return True
119 def process_solution(L):
120 """
121 Data to store in a file. Called on the class, not an instance.
123 Args:
124 L (pySDC.Level.level): the level to log
126 Returns:
127 dict: the time ``t`` at the end of the step and a numpy view ``u`` of ``L.uend``
128 """
129 return {'t': L.time + L.dt, 'u': L.uend.view(np.ndarray)}
131 def format_index(index):
132 """
133 Render the file index for the file name. Called on the class, not an instance.
135 Args:
136 index (int): the index of the file
138 Returns:
139 str: the index padded with zeros to six digits
140 """
141 return f'{index:06d}'
143 def __init__(self):
144 """Check that ``path`` is set and is not a file, and create the directory if it does not exist."""
145 super().__init__()
147 if self.path is None:
148 raise ValueError('Please set a path for logging as the class attribute `LogToFile.path`!')
150 if os.path.isfile(self.path):
151 raise ValueError(
152 f'{self.path!r} is not a valid path to log to because a file of the same name exists. Please supply a directory'
153 )
155 if not os.path.isdir(self.path):
156 os.makedirs(self.path, exist_ok=True)
158 def log_to_file(self, step, level_number, condition, process_solution=None):
159 """
160 Pickle the processed solution of the finest level to the file of the current ``counter`` and increment the
161 counter. Nothing is written on coarser levels or if ``condition`` is false.
163 Args:
164 step (pySDC.Step.step): the current step
165 level_number (int): the current level number
166 condition (bool): whether to write the file
167 process_solution (callable): replaces the class's ``process_solution`` if given
169 Returns:
170 None
171 """
172 if level_number > 0:
173 return None
175 L = step.levels[level_number]
177 if condition:
178 path = self.get_path(self.counter)
180 if process_solution:
181 data = process_solution(L)
182 else:
183 data = type(self).process_solution(L)
185 with open(path, 'wb') as file:
186 pickle.dump(data, file)
187 self.logger.info(f'Stored file {path!r}')
189 type(self).counter += 1
191 def post_step(self, step, level_number):
192 """
193 Write the solution at the end of the step to file if ``logging_condition`` allows it.
195 Args:
196 step (pySDC.Step.step): the current step
197 level_number (int): the current level number
199 Returns:
200 None
201 """
202 super().post_step(step, level_number)
203 L = step.levels[level_number]
204 self.log_to_file(step, level_number, type(self).logging_condition(L))
206 def pre_run(self, step, level_number):
207 """
208 Write the initial conditions to file, stored with the time ``L.time`` at the start of the run. Sets ``L.uend``
209 to the initial conditions for this, since ``process_solution`` reads ``L.uend``; the first end point computed in
210 the run replaces it.
212 Args:
213 step (pySDC.Step.step): the current step
214 level_number (int): the current level number
216 Returns:
217 None
218 """
219 super().pre_run(step, level_number)
220 L = step.levels[level_number]
221 L.uend = L.u[0]
223 def process_solution(L):
224 return {
225 **type(self).process_solution(L),
226 't': L.time,
227 }
229 self.log_to_file(step, level_number, True, process_solution=process_solution)
231 @classmethod
232 def get_path(cls, index):
233 """
234 Get the path to the file with a given index.
236 Args:
237 index (int): the index of the file
239 Returns:
240 str: the path
241 """
242 return f'{cls.path}/{cls.file_name}_{cls.format_index(index)}.pickle'
244 @classmethod
245 def load(cls, index):
246 """
247 Load the data stored in the file with a given index.
249 Args:
250 index (int): the index of the file
252 Returns:
253 dict: the data as returned by ``process_solution`` when it was written
254 """
255 path = cls.get_path(index)
256 with open(path, 'rb') as file:
257 return pickle.load(file)
260class LogToPickleFileAfterXS(LogToPickleFile):
261 r'''
262 Log to file after certain amount of time has passed instead of after every step
263 '''
265 time_increment = 0
266 t_next_log = 0
268 def post_step(self, step, level_number):
269 """
270 Write the solution to file once the end of the step has reached the next logging time and the step is not
271 restarted, then move the next logging time on by ``time_increment``.
273 Args:
274 step (pySDC.Step.step): the current step
275 level_number (int): the current level number
277 Returns:
278 None
279 """
280 L = step.levels[level_number]
282 if self.t_next_log == 0:
283 self.t_next_log = self.time_increment
285 # rounding in the accumulated step times must not skip a logging time
286 if L.time + L.dt >= self.t_next_log - np.finfo(float).eps * 1000 and not step.status.restart:
287 super().post_step(step, level_number)
288 self.t_next_log = max([L.time + L.dt, self.t_next_log]) + self.time_increment
289 else:
290 # skip writing in the parent class, but not the rest of the chain
291 super(LogToPickleFile, self).post_step(step, level_number)
293 def pre_run(self, step, level_number):
294 """
295 Write the initial conditions to file like the parent class, but only if ``logging_condition`` allows it.
297 Args:
298 step (pySDC.Step.step): the current step
299 level_number (int): the current level number
301 Returns:
302 None
303 """
304 # replaces the writing in the parent class, but not the rest of the chain
305 super(LogToPickleFile, self).pre_run(step, level_number)
306 L = step.levels[level_number]
307 L.uend = L.u[0]
309 def process_solution(L):
310 return {
311 **type(self).process_solution(L),
312 't': L.time,
313 }
315 self.log_to_file(step, level_number, type(self).logging_condition(L), process_solution=process_solution)
318class LogToFile(Hooks):
319 """
320 Write the solution every `time_increment` to one `FieldsIO` file set up by the problem, resuming an existing file.
321 """
323 filename = 'myRun.pySDC'
324 time_increment = 0
325 allow_overwriting = False
326 counter = 0 # number of stored time points in the file
328 def __init__(self):
329 """Set up without an output file and allow overwriting in ``FieldsIO`` according to ``allow_overwriting``."""
330 super().__init__()
331 self.outfile = None
332 self.t_next_log = 0
333 FieldsIO.ALLOW_OVERWRITE = self.allow_overwriting
335 def pre_run(self, step, level_number):
336 """
337 Set up the output file on the finest level. If the file exists and the run does not start at :math:`t=0`, it is
338 opened to append to it. Otherwise, the problem creates a new file and the initial conditions are written to it.
340 Args:
341 step (pySDC.Step.step): the current step
342 level_number (int): the current level number
344 Returns:
345 None
346 """
347 super().pre_run(step, level_number)
348 if level_number > 0:
349 return None
350 L = step.levels[level_number]
352 # setup outfile
353 if os.path.isfile(self.filename) and L.time > 0:
354 L.prob.setUpFieldsIO()
355 self.outfile = FieldsIO.fromFile(self.filename)
356 self.counter = len(self.outfile.times)
357 self.logger.info(
358 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}.'
359 )
360 else:
361 self.outfile = L.prob.getOutputFile(self.filename)
362 self.logger.info(f'Set up file {self.filename!r} for writing output.')
364 # write initial conditions
365 if L.time not in self.outfile.times:
366 self.outfile.addField(time=L.time, field=L.prob.processSolutionForOutput(L.u[0]))
367 self.logger.info(f'Written initial conditions at t={L.time:4f} to file')
369 type(self).counter = len(self.outfile.times)
370 self.logger.info(f'Will write to disk every {self.time_increment:.4e} time units')
372 def post_step(self, step, level_number):
373 """
374 Write the solution of the finest level to file once the end of the step has reached the next logging time and
375 the step is not restarted, then move the next logging time on by ``time_increment``.
377 Args:
378 step (pySDC.Step.step): the current step
379 level_number (int): the current level number
381 Returns:
382 None
384 Raises:
385 DataError: if the file already has a solution at this time and ``allow_overwriting`` is not set
386 """
387 super().post_step(step, level_number)
388 if level_number > 0:
389 return None
391 L = step.levels[level_number]
393 if self.t_next_log == 0:
394 self.t_next_log = L.time + self.time_increment
396 # rounding in the accumulated step times must not skip a logging time
397 if L.time + L.dt >= self.t_next_log - np.finfo(float).eps * 1000 and not step.status.restart:
398 value_exists = True in [abs(me - (L.time + L.dt)) < np.finfo(float).eps * 1000 for me in self.outfile.times]
399 if value_exists and not self.allow_overwriting:
400 raise DataError(f'Already have recorded data for time {L.time + L.dt} in this file!')
401 self.outfile.addField(time=L.time + L.dt, field=L.prob.processSolutionForOutput(L.uend))
402 self.logger.info(f'Written solution at t={L.time+L.dt:.4f} to file')
403 self.t_next_log = max([L.time + L.dt, self.t_next_log]) + self.time_increment
404 type(self).counter = len(self.outfile.times)
406 def post_run(self, step, level_number):
407 """
408 Write the final solution of the finest level to file unless the file has it already.
410 Args:
411 step (pySDC.Step.step): the current step
412 level_number (int): the current level number
414 Returns:
415 None
416 """
417 super().post_run(step, level_number)
418 if level_number > 0:
419 return None
421 L = step.levels[level_number]
423 value_exists = True in [abs(me - (L.time + L.dt)) < np.finfo(float).eps * 1000 for me in self.outfile.times]
424 if not value_exists:
425 self.outfile.addField(time=L.time + L.dt, field=L.prob.processSolutionForOutput(L.uend))
426 self.logger.info(f'Written solution at t={L.time+L.dt:.4f} to file')
427 self.t_next_log = max([L.time + L.dt, self.t_next_log]) + self.time_increment
428 type(self).counter = len(self.outfile.times)
430 @classmethod
431 def load(cls, index):
432 """
433 Load a solution from ``filename``.
435 Args:
436 index (int): the index of the time point in the file
438 Returns:
439 dict: the time ``t`` and the solution ``u``
440 """
441 data = {}
442 file = FieldsIO.fromFile(cls.filename)
443 file_entry = file.readField(idx=index)
444 data['u'] = file_entry[1]
445 data['t'] = file_entry[0]
446 return data