Coverage for gamdpy/interactions/planar.py: 66%

1from typing import Callable 

2 

3import numba 

4import numpy as np 

5from numba import cuda 

6 

7from gamdpy import Configuration 

8from . import Interaction 

9from .make_fixed_interactions import make_fixed_interactions 

10 

11class Planar(Interaction): 

12 """ Planar interactions such as smooth walls, gravity, or an electric field. 

13 

14 Consider a plane with the normal vector :math:`{\\bf n}` going though the point :math:`{\\bf p}`. 

15 For a given particle, let :math:`{\\bf r}` be the distance to the nearest point in the plane. 

16 Then the planar for is 

17 

18 .. math:: 

19 

20 {\\bf F} = s(r) {\\bf n} 

21 

22 Where :math:`s(r)=-u'(r)/r` is the force multiplier of a given potential function. 

23 

24 Note: The planer interaction is considered an "external force", 

25 and will not contribute particles scalar energies, virials etc. 

26 

27 Parameters 

28 ---------- 

29 

30 potential : Callable 

31 Potential function for planer interactions 

32 See :func:gamdpy.potential_functions.harmonic_repulsion for an example. 

33 

34 params : list[list[float]] 

35 A list of parameters for each plane type. 

36 Each entry is a list of parameters for the potential function (see above). 

37 

38 indices : list[list[int]] 

39 A list of lists, each containing the indices of a particle involved, and the planer interactions of relevance. 

40 

41 normal_vectors : list[list[float]] 

42 A list of lists, each containing a normal vector for a given plane. 

43 

44 points : list[list[float]] 

45 A list of lists, each containing a point on a given plane. 

46 

47 """ 

48 

49 def __init__(self, 

50 potential: Callable, 

51 params: list[list[float]], 

52 indices: list[list[int]], 

53 normal_vectors: list[list[float]], 

54 points: list[list[float]]): 

55 

56 # User set variables 

57 self.potential = potential 

58 self.params = params 

59 self.indices = indices # list with "particle index", "planer interaction type" 

60 self.normal_vectors = normal_vectors 

61 self.points = points 

62 

63 # Derived variables 

64 self.potential_njit = numba.njit(potential) 

65 self.num_types = len(self.params) 

66 

67 # Set by get_kernal method 

68 self.d_indices = None 

69 self.d_params = None 

70 

71 

72 def get_kernel(self, configuration: Configuration, compute_plan: dict, compute_flags: dict[str,bool]) -> Callable: 

73 """ Get a kernel that implements calculation of the interaction """ 

74 D = configuration.D 

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

76 dist_sq_function = numba.njit(configuration.simbox.get_dist_sq_function()) 

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

78 pot_func = self.potential_njit 

79 

80 def calculator(vectors, scalars, ptype, sim_box, indices, values): 

81 particle = indices[0] 

82 interaction_type = indices[1] 

83 point = values[interaction_type][0:D] # Point on plane 

84 normal_vector = values[interaction_type][D:2 * D] # Normal vector defining plane 

85 potential_params = values[interaction_type][2 * D:] 

86 

87 # Calculating displacement vector to plane 

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

89 dist_sq = dist_sq_dr_function(vectors[r_id][particle], point, sim_box, dr) 

90 dist = numba.float32(0.0) 

91 for k in range(D): 

92 dist += dr[k] * normal_vector[k] 

93 

94 if particle==0 and interaction_type==1: 

95 #print(interaction_type) 

96 #print(point[0],point[1], point[2] ) 

97 #print(normal_vector[0],normal_vector[1],normal_vector[2] ) 

98 #print(dist_sq) 

99 #print(dist, values[interaction_type][-1]) 

100 #print(potential_params[0], potential_params[1], potential_params[2], ) 

101 pass 

102 

103 if abs(dist) < values[interaction_type][-1]: # Last index is the cut-off 

104 u, s, umm = pot_func(dist, potential_params) 

105 #print(particle, vectors[r_id][particle][1], interaction_type, dist, s, normal_vector[0],normal_vector[1],normal_vector[2]) 

106 for k in range(D): 

107 cuda.atomic.add(vectors, (f_id, particle, k), normal_vector[k] * dist * s) # Force 

108 

109 return 

110 

111 return make_fixed_interactions(configuration, calculator, compute_plan, verbose=False) 

112 

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

114 """ Get a tuple with the parameters expected by the associated kernel """ 

115 values = [] 

116 for point, vector, potenetial_params in zip(self.points, self.normal_vectors, self.params): 

117 values.append(point + vector + potenetial_params) 

118 

119 self.d_indices = cuda.to_device(self.indices) 

120 self.d_values = cuda.to_device(values) 

121 return self.d_indices, self.d_values