Coverage for gamdpy/runtime_actions/momentum_reset.py: 63%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda 

5 

6# Abstract Base Class and type annotation 

7from .runtime_action import RuntimeAction 

8from gamdpy import Configuration 

9 

10 

11# Could include flags of dimensions to work on 

12class MomentumReset(RuntimeAction): 

13 """  

14 Runtime action that sets the total momentum of configuration to zero 

15 every `steps_between_reset` time step. 

16 """ 

17 

18 def __init__(self, steps_between_reset: int) -> None: 

19 if type(steps_between_reset) != int or steps_between_reset < 0: 

20 raise ValueError(f'steps_between_momentum_reset ({steps_between_reset}) should be non-negative integer.') 

21 self.steps_between_reset = steps_between_reset 

22 

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

24 pass 

25 

26 def get_params(self, configuration: Configuration, compute_plan: dict) -> tuple: 

27 self.total_momentum = np.zeros(configuration.D+1, dtype=np.float32) # Total mass summed in last index of total_momentum 

28 self.d_total_momentum = cuda.to_device(self.total_momentum) 

29 return (self.d_total_momentum, ) # return parameters as a tuple 

30 

31 def get_prestep_kernel(self, configuration: Configuration, compute_plan: dict): 

32 

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

34 if gridsync: 

35 def kernel(grid, vectors, scalars, r_im, sim_box, step, momentum_reset_params): 

36 pass 

37 return 

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

39 else: 

40 def kernel(grid, vectors, scalars, r_im, sim_box, step, momentum_reset_params): 

41 pass 

42 return kernel 

43 

44 def get_poststep_kernel(self, configuration: Configuration, compute_plan: dict): 

45 

46 # Unpack parameters from configuration and compute_plan 

47 D, num_part = configuration.D, configuration.N 

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

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

50 

51 # Unpack indices for vectors and scalars to be compiled into kernel 

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

53 m_id = configuration.sid['m'] 

54 

55 steps_between_reset = self.steps_between_reset 

56 

57 def zero_momentum(cm_velocity): 

58 global_id, my_t = cuda.grid(2) 

59 if global_id == 0 and my_t == 0: 

60 for k in range(D+1): 

61 cm_velocity[k] = np.float32(0.) 

62 return 

63 

64 def sum_momentum(vectors, scalars, cm_velocity): 

65 global_id, my_t = cuda.grid(2) 

66 if global_id < num_part and my_t == 0: 

67 my_m = scalars[global_id][m_id] 

68 for k in range(D): 

69 cuda.atomic.add(cm_velocity, k, my_m * vectors[v_id][global_id][k]) 

70 cuda.atomic.add(cm_velocity, D, my_m) # Total mass summed in last index of cm_velocity 

71 return 

72 

73 def shift_velocities(vectors, cm_velocity): 

74 global_id, my_t = cuda.grid(2) 

75 if global_id < num_part and my_t == 0: 

76 for k in range(D): 

77 vectors[v_id][global_id,k] -= cm_velocity[k] / cm_velocity[D] 

78 return 

79 

80 zero_momentum = cuda.jit(device=gridsync)(zero_momentum) 

81 sum_momentum = cuda.jit(device=gridsync)(sum_momentum) 

82 shift_velocities = cuda.jit(device=gridsync)(shift_velocities) 

83 

84 if gridsync: 

85 def kernel(grid, vectors, scalars, r_im, sim_box, step, momentum_reset_params): 

86 cm_velocity, = momentum_reset_params 

87 if step%steps_between_reset == 0: 

88 zero_momentum(cm_velocity) 

89 grid.sync() 

90 sum_momentum(vectors, scalars, cm_velocity) 

91 grid.sync() 

92 shift_velocities(vectors, cm_velocity) 

93 return 

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

95 else: 

96 def kernel(grid, vectors, scalars, r_im, sim_box, step, momentum_reset_params): 

97 cm_velocity, = momentum_reset_params 

98 if step%steps_between_reset == 0: 

99 zero_momentum[1, 1](cm_velocity) 

100 sum_momentum[num_blocks, (pb, 1)](vectors, scalars, cm_velocity) 

101 shift_velocities[num_blocks, (pb, 1)](vectors, cm_velocity) 

102 return 

103 return kernel 

104