Coverage for gamdpy/runtime_actions/trajectory_saver.py: 66%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda, config 

5 

6from .runtime_action import RuntimeAction 

7 

8 

9class TrajectorySaver(RuntimeAction): 

10 """  

11 Runtime action for saving configurations during timeblock. 

12 Does logarithmic saving. 

13 """ 

14 

15 def __init__(self, include_simbox=False, verbose=False, compression="gzip", compression_opts=4) -> None: 

16 

17 self.include_simbox = include_simbox 

18 self.num_vectors = 2 # 'r' and 'r_im' (for now!) 

19 self.compression = compression 

20 if self.compression == 'gzip': 

21 self.compression_opts = compression_opts 

22 else: 

23 self.compression_opts = None 

24 #self.sid = {"r":0, "r_im":1} 

25 

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

27 self.configuration = configuration 

28 

29 if type(num_timeblocks) != int or num_timeblocks < 0: 

30 raise ValueError(f'num_timeblocks ({num_timeblocks}) should be non-negative integer.') 

31 self.num_timeblocks = num_timeblocks 

32 

33 if type(steps_per_timeblock) != int or steps_per_timeblock < 0: 

34 raise ValueError(f'steps_per_timeblock ({steps_per_timeblock}) should be non-negative integer.') 

35 self.steps_per_timeblock = steps_per_timeblock 

36 

37 self.conf_per_block = int(math.log2(self.steps_per_timeblock)) + 2 # Should be user controllable 

38 

39 # Setup output 

40 if verbose: 

41 print(f'Storing results in memory. Expected footprint {self.num_timeblocks * self.conf_per_block * self.num_vectors * self.configuration.N * self.configuration.D * 4 / 1024 / 1024:.2f} MB.') 

42 

43 if 'trajectory_saver' in output.keys(): 

44 del output['trajectory_saver'] 

45 output.create_group('trajectory_saver') 

46 

47 # Compression has a different syntax depending if is gzip or not because gzip can have also a compression_opts 

48 # it is possible to use compression=None for not compressing the data 

49 output.create_dataset('trajectory_saver/positions', 

50 shape=(self.num_timeblocks, self.conf_per_block, self.configuration.N, self.configuration.D), 

51 chunks=(1, 1, self.configuration.N, self.configuration.D), 

52 dtype=np.float32, compression=self.compression, compression_opts=self.compression_opts) 

53 output.create_dataset('trajectory_saver/images', 

54 shape=(self.num_timeblocks, self.conf_per_block, self.configuration.N, self.configuration.D), 

55 chunks=(1, 1, self.configuration.N, self.configuration.D), 

56 dtype=np.int32, compression=self.compression, compression_opts=self.compression_opts) 

57 output['trajectory_saver'].attrs['compression_info'] = f"{self.compression} with opts {self.compression_opts}" 

58 

59 #output.attrs['vectors_names'] = list(self.sid.keys()) 

60 if self.include_simbox: 

61 if 'sim_box' in output['trajectory_saver'].keys(): 

62 del output['trajectory_saver/sim_box'] 

63 output.create_dataset('trajectory_saver/sim_box', 

64 shape=(self.num_timeblocks, self.conf_per_block, self.configuration.simbox.len_sim_box_data)) 

65 

66 flag = config.CUDA_LOW_OCCUPANCY_WARNINGS 

67 config.CUDA_LOW_OCCUPANCY_WARNINGS = False 

68 self.zero_kernel = self.make_zero_kernel() 

69 config.CUDA_LOW_OCCUPANCY_WARNINGS = flag 

70 

71 def get_params(self, configuration, compute_plan): 

72 self.conf_array = np.zeros((self.conf_per_block, self.num_vectors, self.configuration.N, self.configuration.D), 

73 dtype=np.float32) 

74 self.d_conf_array = cuda.to_device(self.conf_array) 

75 

76 if self.include_simbox: 

77 self.sim_box_output_array = np.zeros((self.conf_per_block, self.configuration.simbox.len_sim_box_data), dtype=np.float32) 

78 self.d_sim_box_output_array = cuda.to_device(self.sim_box_output_array) 

79 return (self.d_conf_array, self.d_sim_box_output_array) 

80 else: 

81 return (self.d_conf_array,) 

82 

83 def make_zero_kernel(self): 

84 # Unpack parameters from configuration and compute_plan 

85 D, num_part = self.configuration.D, self.configuration.N 

86 pb = 32 

87 num_blocks = (num_part - 1) // pb + 1 

88 

89 def zero_kernel(array): 

90 Nx, Ny, Nz, Nw = array.shape 

91 global_id = cuda.grid(1) 

92 

93 if global_id < Nz: # particles 

94 for i in range(Nx): 

95 for j in range(Ny): 

96 for k in range(Nw): 

97 array[i, j, global_id, k] = numba.float32(0.0) 

98 

99 zero_kernel = cuda.jit(zero_kernel) 

100 return zero_kernel[num_blocks, pb] 

101 

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

103 data = self.d_conf_array.copy_to_host() 

104 # note that d_conf_array has dimensions (self.conf_per_block, 2, self.configuration.N, self.configuration.D) 

105 output_reference['trajectory_saver/positions'][timeblock], output_reference['trajectory_saver/images'][timeblock] = data[:, 0], data[:, 1] 

106 #output['trajectory_saver'][block, :] = self.d_conf_array.copy_to_host() 

107 if self.include_simbox: 

108 output_reference['trajectory_saver/sim_box'][timeblock, :] = self.d_sim_box_output_array.copy_to_host() 

109 self.zero_kernel(self.d_conf_array) 

110 

111 def get_poststep_kernel(self, configuration, compute_plan, verbose=False): 

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

113 if gridsync: 

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

115 pass 

116 return 

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

118 else: 

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

120 pass 

121 return kernel 

122 

123 

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

125 # Unpack parameters from configuration and compute_plan 

126 D, num_part = configuration.D, configuration.N 

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

128 num_blocks = (num_part - 1) // pb + 1 

129 sim_box_array_length = configuration.simbox.len_sim_box_data 

130 include_simbox = self.include_simbox 

131 

132 # Unpack indices for scalars to be compiled into kernel  

133 r_id, = [configuration.vectors.indices[key] for key in ['r', ]] 

134 

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

136 if include_simbox: 

137 conf_array, sim_box_output_array = conf_saver_params 

138 else: 

139 conf_array, = conf_saver_params 

140 

141 Flag = False 

142 if step == 0: 

143 Flag = True 

144 save_index = 0 

145 else: 

146 b = np.int32(math.log2(np.float32(step))) 

147 c = 2 ** b 

148 if step == c: 

149 Flag = True 

150 save_index = b + 1 

151 

152 if Flag: 

153 global_id, my_t = cuda.grid(2) 

154 if global_id < num_part and my_t == 0: 

155 for k in range(D): 

156 conf_array[save_index, 0, global_id, k] = vectors[r_id][global_id, k] 

157 conf_array[save_index, 1, global_id, k] = np.float32(r_im[global_id, k]) 

158 if include_simbox and global_id == 0: 

159 for k in range(sim_box_array_length): 

160 sim_box_output_array[save_index, k] = sim_box[k] 

161 return 

162 

163 kernel = cuda.jit(device=gridsync)(kernel) 

164 

165 if gridsync: 

166 return kernel # return device function 

167 else: 

168 return kernel[num_blocks, (pb, 1)] # return kernel, incl. launch parameters