Coverage for gamdpy/runtime_actions/restart_saver.py: 34%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda, config 

5import h5py 

6 

7from .runtime_action import RuntimeAction 

8 

9class RestartSaver(RuntimeAction): 

10 """  

11 Runtime action for saving restarts, ie. the current configuration at beginning of every timebock . 

12 """ 

13 

14 def __init__(self, timeblocks_between_restarts=1) -> None: 

15 # Later: Give user influence on what and how often is saved 

16 self.timeblocks_between_restarts = timeblocks_between_restarts 

17 

18 def setup(self, configuration, num_timeblocks:int, steps_per_timeblock:int, output, verbose=False) -> None: 

19 

20 self.configuration = configuration 

21 

22 # Setup output 

23 if 'restarts' in output.keys(): 

24 del output['restarts'] 

25 grp = output.create_group('restarts') 

26 grp.attrs['timeblocks_between_restarts'] = self.timeblocks_between_restarts 

27 

28 

29 def get_params(self, configuration, compute_plan): 

30 return (0,) 

31 

32 def initialize_before_timeblock(self, timeblock: int, output_reference): 

33 self.configuration.save(output=output_reference, group_name=f"/restarts/restart{timeblock:04d}", mode="w", include_topology=True) 

34 

35 

36 def update_at_end_of_timeblock(self, timeblock: int, output_reference): 

37 pass 

38 

39 def get_prestep_kernel(self, configuration, compute_plan, verbose=False): 

40 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']] 

41 if gridsync: 

42 def kernel(grid, vectors, scalars, r_im, sim_box, step, conf_saver_params): 

43 pass 

44 return 

45 return cuda.jit(device=gridsync)(kernel) 

46 else: 

47 def kernel(grid, vectors, scalars, r_im, sim_box, step, conf_saver_params): 

48 pass 

49 return kernel 

50 

51 def get_poststep_kernel(self, configuration, compute_plan): 

52 # Unpack parameters from configuration and compute_plan 

53 

54 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']] 

55 if gridsync: 

56 def kernel(grid, vectors, scalars, r_im, sim_box, step, conf_saver_params): 

57 pass 

58 return 

59 return cuda.jit(device=gridsync)(kernel) 

60 else: 

61 def kernel(grid, vectors, scalars, r_im, sim_box, step, conf_saver_params): 

62 pass 

63 return kernel