Coverage for gamdpy/runtime_actions/stress_saver.py: 69%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda, config 

5 

6from .runtime_action import RuntimeAction 

7 

8 

9class StressSaver(RuntimeAction): 

10 """ Runtime action for saving stress tensor(a D x D matrix) during a timeblock 

11 every `steps_between_output` time steps. 

12 """ 

13 def __init__(self, steps_between_output:int = 16, compute_flags = None, verbose=False) -> None: 

14 if type(steps_between_output) != int or steps_between_output < 0: 

15 raise ValueError(f'steps_between_output ({steps_between_output}) should be non-negative integer.') 

16 self.steps_between_output = steps_between_output 

17 self.compute_flags = compute_flags 

18 self.verbose = verbose 

19 

20 def get_compute_flags(self): 

21 return self.compute_flags 

22 

23 

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

25 

26 self.configuration = configuration 

27 D = configuration.D 

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 if self.steps_between_output >= steps_per_timeblock: 

38 raise ValueError(f'scalar_output ({self.steps_between_output}) must be less than steps_per_timeblock ({steps_per_timeblock})') 

39 

40 # per block saving of stress tensor 

41 if not configuration.compute_flags['stresses']: 

42 raise RuntimeError('stresses not set in compute_flags') 

43 

44 self.stress_saves_per_block = self.steps_per_timeblock//self.steps_between_output 

45 

46 # Setup output 

47 shape = (self.num_timeblocks, self.stress_saves_per_block, D, D) 

48 if 'stress_saver' in output.keys(): 

49 del output['stress_saver'] 

50 grp = output.create_group('stress_saver') 

51 output.create_dataset('stress_saver/stress_tensor', shape=shape, 

52 chunks=(1, self.stress_saves_per_block, D, D), dtype=np.float32) 

53 grp.attrs['steps_between_output'] = self.steps_between_output 

54 

55 flag = config.CUDA_LOW_OCCUPANCY_WARNINGS 

56 config.CUDA_LOW_OCCUPANCY_WARNINGS = False 

57 self.zero_kernel = self.make_zero_kernel_3() 

58 config.CUDA_LOW_OCCUPANCY_WARNINGS = flag 

59 

60 

61 

62 def make_zero_kernel_3(self): 

63 """ Returns a kernel that can zero an array with three axes """ 

64 def zero_kernel(array): 

65 Nx, Ny, Nz = array.shape 

66 #i, j = cuda.grid(2) # doing simple 1 thread kernel for now ... 

67 for i in range(Nx): 

68 for j in range(Ny): 

69 for k in range(Nz): 

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

71 

72 zero_kernel = cuda.jit(zero_kernel) 

73 return zero_kernel[1,1] 

74 

75 

76 def get_params(self, configuration, compute_plan): 

77 D = configuration.D 

78 self.output_array = np.zeros((self.stress_saves_per_block, D, D), dtype=np.float32) 

79 self.d_output_array = cuda.to_device(self.output_array) 

80 self.params = (self.steps_between_output, self.d_output_array) 

81 return self.params 

82 

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

84 self.zero_kernel(self.d_output_array) 

85 

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

87 volume = self.configuration.get_volume() 

88 output_reference['stress_saver/stress_tensor'][timeblock, :] = self.d_output_array.copy_to_host() / volume 

89 

90 

91 def get_prestep_kernel(self, configuration, compute_plan): 

92 # Unpack parameters from configuration and compute_plan 

93 D, num_part = configuration.D, configuration.N 

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

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

96 

97 

98 m_id = configuration.sid['m'] 

99 v_id = configuration.vectors.indices['v'] 

100 sx_id = configuration.vectors.indices['sx'] 

101 sy_id = configuration.vectors.indices['sy'] 

102 if D > 2: 

103 sz_id = configuration.vectors.indices['sz'] 

104 if D > 3: 

105 sw_id = configuration.vectors.indices['sw'] 

106 

107 #volume_function = numba.njit(configuration.simbox.get_volume_function()) 

108 

109 def kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_action_params): 

110 """  

111 """ 

112 steps_between_output, output_array = runtime_action_params # Needs to be compatible with get_params above 

113 if step%steps_between_output==0: 

114 save_index = step//steps_between_output 

115 if save_index < output_array.shape[0]: 

116 global_id, my_t = cuda.grid(2) 

117 if global_id < num_part and my_t == 0: 

118 my_m = scalars[global_id][m_id] 

119 for k in range(D): 

120 cuda.atomic.add(output_array, (save_index, 0, k), vectors[sx_id][global_id][k] - 

121 my_m * vectors[v_id][global_id][0]*vectors[v_id][global_id][k]) 

122 

123 cuda.atomic.add(output_array, (save_index, 1, k), vectors[sy_id][global_id][k] - 

124 my_m * vectors[v_id][global_id][1]*vectors[v_id][global_id][k]) 

125 if D > 2: 

126 cuda.atomic.add(output_array, (save_index, 2, k), vectors[sz_id][global_id][k] - 

127 my_m * vectors[v_id][global_id][2]*vectors[v_id][global_id][k]) 

128 if D > 3: 

129 cuda.atomic.add(output_array, (save_index, 3, k), vectors[sw_id][global_id][k] - 

130 my_m * vectors[v_id][global_id][3]*vectors[v_id][global_id][k]) 

131 return 

132 

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

134 

135 if gridsync: 

136 return kernel # return device function 

137 else: 

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

139 

140 #### CONSIDER USING POSTSTEP TO GET CORRECT VELOCITIES ############### 

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

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

143 if gridsync: 

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

145 pass 

146 return 

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

148 else: 

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

150 pass 

151 return kernel 

152 

153 # Class functions to read data 

154 

155 def extract(h5file, first_block=0, last_block=None, subsample=1): 

156 

157 #stress_data = h5file['stress_saver/stress_tensor'] 

158 h5grp = h5file['stress_saver'] 

159 nblocks, per_block, D, D2 = h5grp['stress_tensor'].shape 

160 assert D == D2 

161 final_rows = (nblocks-first_block) * per_block 

162 return h5grp['stress_tensor'][first_block:last_block,:,:, :].reshape(final_rows, D, D)[::subsample] 

163 

164 def get_times(h5file, first_block=0, last_block=-1, reset_time=True, subsample=1): 

165 num_timeblock, saves_per_timeblock = h5file['stress_saver']['stress_tensor'][first_block:last_block,:,0,0].shape 

166 times_array = np.arange(0,num_timeblock*saves_per_timeblock, step=subsample)*h5file.attrs['dt'] 

167 return times_array 

168 

169def extract_stress_tensor(data, first_block=0, D=3): 

170 """ Extracts stress tensor data from simulation output. 

171 

172 Parameters 

173 ---------- 

174 

175 data : dict 

176 Output from a Simulation object. 

177 

178 

179 

180 first_block : int 

181 Index of the first timeblock to extract data from. 

182 

183 D : int 

184 Dimension of the simulation. 

185 

186 Returns 

187 ------- 

188 

189 array 

190 three-dimensional numpy array containing the extracted tensor data. 

191 

192 """ 

193 stress_data = data['stress_saver/stress_tensor'] 

194 nblocks, per_block, _, _ = stress_data.shape 

195 final_rows = (nblocks-first_block) * per_block 

196 return stress_data[first_block:,:,:].reshape(final_rows, D, D)