Coverage for gamdpy/runtime_actions/scalar_saver.py: 61%
153 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1import numpy as np
2import numba
3import math
4from numba import cuda, config
5import h5py
7from .runtime_action import RuntimeAction
9class ScalarSaver(RuntimeAction):
10 """
11 Runtime action for saving scalar data (such as thermodynamic properties) during a timeblock
12 every `steps_between_output` time steps.
13 """
15 def __init__(self, steps_between_output:int = 16, compute_flags = None, verbose=False, compression="gzip", compression_opts=4) -> None:
17 if type(steps_between_output) != int or steps_between_output < 0:
18 raise ValueError(f'steps_between_output ({steps_between_output}) should be non-negative integer.')
19 self.steps_between_output = steps_between_output
21 self.compute_flags = compute_flags
22 self.compression = compression
23 if self.compression == 'gzip':
24 self.compression_opts = compression_opts
25 else:
26 self.compression_opts = None
28 def get_compute_flags(self):
29 return self.compute_flags
31 def setup(self, configuration, num_timeblocks:int, steps_per_timeblock:int, output, verbose=False) -> None:
33 self.configuration = configuration
35 if type(num_timeblocks) != int or num_timeblocks < 0:
36 raise ValueError(f'num_timeblocks ({num_timeblocks}) should be non-negative integer.')
37 self.num_timeblocks = num_timeblocks
39 if type(steps_per_timeblock) != int or steps_per_timeblock < 0:
40 raise ValueError(f'steps_per_timeblock ({steps_per_timeblock}) should be non-negative integer.')
41 self.steps_per_timeblock = steps_per_timeblock
43 if self.steps_between_output >= steps_per_timeblock:
44 raise ValueError(f'scalar_output ({self.steps_between_output}) must be less than steps_per_timeblock ({steps_per_timeblock})')
46 # per block saving of scalars
47 compute_flags = configuration.compute_flags
48 self.num_scalars = 0
49 sid_list = ['U', 'W', 'lapU', 'K', 'Fsq', 'Vol']
50 self.sid = {}
51 for item in sid_list:
52 if compute_flags[item]:
53 self.sid[item] = self.num_scalars
54 self.num_scalars += 1
56 if compute_flags['Ptot']:
57 self.sid['Px'] = self.num_scalars
58 self.sid['Py'] = self.num_scalars + 1
59 self.sid['Pz'] = self.num_scalars + 2
60 self.num_scalars += 3
62 if compute_flags['stresses']:
63 self.sid['Sxy'] = self.num_scalars
64 self.num_scalars += 1
66 self.scalar_saves_per_block = self.steps_per_timeblock//self.steps_between_output
68 # Setup output
69 shape = (self.num_timeblocks, self.scalar_saves_per_block, self.num_scalars)
70 if 'scalar_saver' in output.keys():
71 del output['scalar_saver']
72 output.create_group('scalar_saver')
74 # Compression has a different syntax depending if is gzip or not because gzip can have also a compression_opts
75 # it is possible to use compression=None for not compressing the data
76 output.create_dataset('scalar_saver/scalars', shape=shape,
77 chunks=(1, self.scalar_saves_per_block, self.num_scalars),
78 dtype=np.float32, compression=self.compression, compression_opts=self.compression_opts)
79 output['scalar_saver'].attrs['compression_info'] = f"{self.compression} with opts {self.compression_opts}"
81 output['scalar_saver'].attrs['steps_between_output'] = self.steps_between_output
82 output['scalar_saver'].attrs['scalar_names'] = list(self.sid.keys())
84 flag = config.CUDA_LOW_OCCUPANCY_WARNINGS
85 config.CUDA_LOW_OCCUPANCY_WARNINGS = False
86 self.zero_kernel = self.make_zero_kernel()
87 config.CUDA_LOW_OCCUPANCY_WARNINGS = flag
89 def make_zero_kernel(self):
91 def zero_kernel(array):
92 Nx, Ny = array.shape
93 #i, j = cuda.grid(2) # doing simple 1 thread kernel for now ...
94 for i in range(Nx):
95 for j in range(Ny):
96 array[i,j] = numba.float32(0.0)
98 zero_kernel = cuda.jit(zero_kernel)
99 return zero_kernel[1,1]
101 def get_params(self, configuration, compute_plan):
103 self.output_array = np.zeros((self.scalar_saves_per_block, self.num_scalars), dtype=np.float32)
104 self.d_output_array = cuda.to_device(self.output_array)
105 self.params = (self.steps_between_output, self.d_output_array)
106 return self.params
108 def initialize_before_timeblock(self, timeblock: int, output_reference):
109 self.zero_kernel(self.d_output_array)
111 def update_at_end_of_timeblock(self, timeblock: int, output_reference):
112 output_reference['scalar_saver/scalars'][timeblock, :] = self.d_output_array.copy_to_host()
114 def get_prestep_kernel(self, configuration, compute_plan, verbose=False):
115 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
116 if gridsync:
117 def kernel(grid, vectors, scalars, r_im, sim_box, step, conf_saver_params):
118 pass
119 return
120 return cuda.jit(device=gridsync)(kernel)
121 else:
122 def kernel(grid, vectors, scalars, r_im, sim_box, step, conf_saver_params):
123 pass
124 return kernel
126 def get_poststep_kernel(self, configuration, compute_plan):
127 # Unpack parameters from configuration and compute_plan
128 D, num_part = configuration.D, configuration.N
129 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']]
130 num_blocks = (num_part - 1) // pb + 1
132 # Below is three alternatives of how to unpack parameters for use in the kernel
134# Alternative 1. Requires renaming of some variables
135# print(configuration.compute_flags)
136# for key in configuration.compute_flags.keys():
137# if key == "stresses" and configuration.compute_flags[key]:
138# sx_id = configuration.vectors.indices['sx']
139# elif configuration.compute_flags[key]:
140# print(key, configuration.compute_flags[key], f'{key}_id', self.sid[key], f'compute_{key}')
141# globals()[f'compute_{key}'] = configuration.compute_flags[key]
142# globals()[f'{key}_id'] = self.sid[key]
143# print(compute_U)
146 # Alternative 2. 'None' used as index for scalars that are not there, and therefore shouldn't be used (kernel throws an getitem error)
147 unpacked_compute_flags = [configuration.compute_flags[key] for key in ['U', 'W', 'lapU', 'Fsq', 'K', 'Vol', 'Ptot', 'stresses']]
148 compute_u, compute_w, compute_lap, compute_fsq, compute_k, compute_vol, compute_Ptot, compute_stresses = unpacked_compute_flags
150 unpacked_scalar_indicies = [self.sid.get(key, None) for key in ['U', 'W', 'lapU', 'Fsq', 'K', 'Vol', 'Px', 'Py', 'Pz', 'Sxy']]
151 u_id, w_id, lap_id, fsq_id, k_id, vol_id, Px_id, Py_id, Pz_id, Sxy_id = unpacked_scalar_indicies
153 # Original
154 #compute_u = configuration.compute_flags['U']
155 #compute_w = configuration.compute_flags['W']
156 #compute_lap = configuration.compute_flags['lapU']
157 #compute_fsq = configuration.compute_flags['Fsq']
158 #compute_k = configuration.compute_flags['K']
159 #compute_vol = configuration.compute_flags['Vol']
160 #compute_Ptot = configuration.compute_flags['Ptot']
161 #compute_stresses = configuration.compute_flags['stresses']
164 # Unpack indices for scalars to be compiled into kernel
165 #if compute_u:
166 # u_id = self.sid['U']
168 #if compute_k:
169 # k_id = self.sid['K']
170 #if compute_w:
171 # w_id = self.sid['W']
173 #if compute_fsq:
174 # fsq_id = self.sid['Fsq']
176 #if compute_lap:
177 # lap_id = self.sid['lapU']
179 #if compute_vol:
180 # vol_id = self.sid['Vol']
182 #if compute_Ptot:
183 # Px_id = self.sid['Px']
184 # Py_id = self.sid['Py']
185 # Pz_id = self.sid['Pz']
187 #if compute_stresses:
188 # Sxy_id = self.sid['Sxy']
190 m_id = configuration.sid['m']
193 v_id = configuration.vectors.indices['v']
194 if compute_stresses:
195 sx_id = configuration.vectors.indices['sx']
197 volume_function = numba.njit(configuration.simbox.get_volume_function())
199 def kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_action_params):
200 """
201 """
202 steps_between_output, output_array = runtime_action_params # Needs to be compatible with get_params above
203 if step%steps_between_output==0:
204 save_index = step//steps_between_output
206 global_id, my_t = cuda.grid(2)
207 if global_id < num_part and my_t == 0:
208 if compute_u:
209 cuda.atomic.add(output_array, (save_index, u_id), scalars[global_id][u_id]) # Potential energy
210 if compute_w:
211 cuda.atomic.add(output_array, (save_index, w_id), scalars[global_id][w_id]) # Virial
212 if compute_lap:
213 cuda.atomic.add(output_array, (save_index, lap_id), scalars[global_id][lap_id]) # Laplace
214 if compute_fsq:
215 cuda.atomic.add(output_array, (save_index, fsq_id), scalars[global_id][fsq_id]) # F**2
216 if compute_k:
217 cuda.atomic.add(output_array, (save_index, k_id), scalars[global_id][k_id]) # Kinetic energy
219 # Contribution to total momentum
220 if compute_Ptot or compute_stresses:
221 my_m = scalars[global_id][m_id]
222 if compute_Ptot:
223 cuda.atomic.add(output_array, (save_index, Px_id), my_m*vectors[v_id][global_id][0])
224 cuda.atomic.add(output_array, (save_index, Py_id), my_m*vectors[v_id][global_id][1])
225 cuda.atomic.add(output_array, (save_index, Pz_id), my_m*vectors[v_id][global_id][2])
227 if compute_stresses:
228 # XY component of stress only for now
229 cuda.atomic.add(output_array, (save_index, Sxy_id), vectors[sx_id][global_id][1] -
230 my_m * vectors[v_id][global_id][0]*vectors[v_id][global_id][1])
232 if compute_vol and global_id == 0 and my_t == 0:
233 output_array[save_index][vol_id] = volume_function(sim_box)
235 return
237 kernel = cuda.jit(device=gridsync)(kernel)
239 if gridsync:
240 return kernel # return device function
241 else:
242 return kernel[num_blocks, (pb, 1)] # return kernel, incl. launch parameters
244 # Class functions to read data
246 def info(h5file):
247 h5grp = h5file['scalar_saver']
248 str = f"\tscalar_names: {h5grp.attrs['scalar_names']}"
249 str += f"\n\tscalars, shape: {h5grp['scalars'].shape}, dtype: {h5grp['scalars'].dtype}"
250 return str
252 def columns(h5file):
253 return list(h5file['scalar_saver'].attrs['scalar_names'])
255 def extract(h5file, columns, per_particle=True, first_block=0, last_block=None, subsample=1, function=None):
256 _, N, D = h5file['initial_configuration']['vectors'].shape
258 h5grp = h5file['scalar_saver']
259 scalar_names = list(h5grp.attrs['scalar_names'])
261 output = []
262 for column in columns:
263 index = scalar_names.index(column)
264 data = np.ravel(h5grp['scalars'][first_block:last_block,:,index])[::subsample]
265 if per_particle:
266 data /= N
268 if function:
269 data = function(data)
271 output.append(data)
272 return output
274 def get_times(h5file, first_block=0, last_block=-1, reset_time=True, subsample=1):
275 num_timeblock, saves_per_timeblock = h5file['scalar_saver']['scalars'][first_block:last_block,:,0].shape
276 times_array = np.arange(0,num_timeblock*saves_per_timeblock, step=subsample)*h5file.attrs['dt']
277 return times_array