Coverage for gamdpy/interactions/tether.py: 81%

1 

2import numba 

3import numpy as np 

4from numba import cuda 

5 

6from .make_fixed_interactions import make_fixed_interactions # tether is an example of 'fixed' interactions 

7 

8# Abstract Base Class and type annotation 

9from .interaction import Interaction 

10from gamdpy import Configuration 

11 

12 

13class Tether(Interaction): 

14 """ Connect particles to anchor-points in space with a harmonic spring force.  

15  

16 Parameters 

17 ---------- 

18 Points and spring constants are defined using either  

19 (a) A list of particle indices to be tethered and associated list of spring constants 

20 (b) A list of particle types to be tethered and associated list of spring constants 

21 

22 See examples/tethered_particles.py 

23 """ 

24 

25 def __init__(self): 

26 self.anchor_points_set = False 

27 

28 def set_anchor_points_from_lists(self, particle_indices, spring_constants, configuration): 

29 """ Set anchor points and spring constants for tethered particles. 

30  

31 Parameters 

32 ---------- 

33 

34 particle_indices : list of int 

35 List of particle indices to be tethered. 

36  

37 spring_constants : list of float 

38 List of spring constants. 

39 

40 configuration : rumd.Configuration 

41 Configuration object containing particle thether positions. 

42  

43 """ 

44 

45 nsprings, nparticles = len(spring_constants), len(particle_indices) 

46 

47 if nsprings != nparticles: 

48 raise ValueError("Each particle must have exactly one spring connection - array must be same length"); 

49 

50 indices, tether_params = [], [] 

51 for n in range(nparticles): 

52 indices.append([n, particle_indices[n]]) 

53 pos = configuration['r'][particle_indices[n]] 

54 tether_params.append( [pos[0], pos[1], pos[2], spring_constants[n]] ) 

55 

56 self.tether_params = np.array(tether_params, dtype=np.float32) 

57 self.indices = np.array(indices, dtype=np.int32) 

58 

59 self.anchor_points_set = True 

60 

61 

62 def set_anchor_points_from_types(self, particle_types, spring_constants, configuration): 

63 

64 nsprings, ntypes = len(spring_constants), len(particle_types) 

65 

66 if ntypes != nsprings: 

67 raise ValueError("Each type must have exactly one spring connection - arrays must be same length") 

68 

69 indices, tether_params, counter = [], [], 0 

70 for n in range(configuration.N): 

71 for m in range(ntypes): 

72 if configuration.ptype[n]==particle_types[m]: 

73 indices.append([counter, n]) 

74 pos = configuration['r'][n] 

75 tether_params.append( [pos[0], pos[1], pos[2], spring_constants[m]] ) 

76 counter = counter + 1 

77 break 

78 

79 self.tether_params = np.array(tether_params, dtype=np.float32) 

80 self.indices = np.array(indices, dtype=np.int32) 

81 

82 self.anchor_points_set = True 

83 

84 def get_params(self, configuration: Configuration, compute_plan: dict, verbose=False) -> tuple: 

85 if self.anchor_points_set == False: 

86 raise ValueError("Anchor points not defined") 

87 

88 self.d_pindices = cuda.to_device(self.indices) 

89 self.d_tether_params = cuda.to_device(self.tether_params); 

90 

91 return (self.d_pindices, self.d_tether_params) 

92 

93 def get_kernel(self, configuration: Configuration, compute_plan: dict, compute_flags: dict[str,bool], verbose=False): 

94 # Unpack parameters from configuration and compute_plan 

95 D, N = configuration.D, configuration.N 

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

97 num_blocks = (N - 1) // pb + 1 

98 

99 compute_u = compute_flags['U'] 

100 # Note w, lap, stresses not relevant here 

101 

102 r_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'f']] 

103 

104 if compute_u: 

105 u_id = configuration.sid['U'] 

106 

107 dist_sq_dr_function = numba.njit(configuration.simbox.get_dist_sq_dr_function()) 

108 

109 def tether_calculator(vectors, scalars, ptype, sim_box, indices, values): 

110 

111 dr = cuda.local.array(shape=D,dtype=numba.float32) 

112 dist_sq = dist_sq_dr_function(values[indices[0]][:D], vectors[r_id][indices[1]], sim_box, dr) 

113 

114 spring = values[indices[0]][3] 

115 

116 f=vectors[f_id][indices[1]]; 

117 

118 for k in range(D): 

119 f[k] = f[k] + dr[k]*spring 

120 

121 Epot = numba.float32(0.5)*spring*dist_sq 

122 # if compute_u: 

123 # What's going on here? Why the atomic add? 

124 cuda.atomic.add(scalars, (indices[0], u_id), Epot) 

125 

126 return 

127 

128 

129 

130 return make_fixed_interactions(configuration, tether_calculator, compute_plan, verbose=False) 

131