Coverage for gamdpy/interactions/planar_interactions.py: 100%

1import numpy as np 

2import numba 

3import math 

4from numba import cuda 

5from .make_fixed_interactions import \ 

6 make_fixed_interactions # planar interactions is an example of 'fixed' interactions 

7 

8 

9######################################################################### 

10######## Planar interactions (smmoth walls, gravity, electric field) 

11######################################################################### 

12 

13def make_planar_calculator(configuration, potential_function) -> callable: 

14 """ Returns a function that calculates a planar interaction for particles 

15 

16 This function is used to create a planar interaction such as a smooth wall, gravity or an electric field. 

17 

18 """ 

19 

20 D = configuration.D 

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

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

23 

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

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

26 u_id, w_id, lap_id = [configuration.sid[key] for key in ['U', 'W', 'lapU']] 

27 

28 def planar_calculator(vectors, scalars, ptype, sim_box, indices, values): # pragma: no cover 

29 particle = indices[0] 

30 interaction_type = indices[1] 

31 point = values[interaction_type][0:D] # Point in wall 

32 normal_vector = values[interaction_type][D:2 * D] # Normal vector defining plane of wall 

33 

34 # Calculating full D-dim displacement vector to avoid worrying about new sim_box types in future 

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

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

37 dist = numba.float32(0.0) 

38 for k in range(D): 

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

40 if dist < values[interaction_type][-1]: # Last index is the cut-off 

41 u, s, umm = potential_function(abs(dist), 

42 values[interaction_type][2 * D:]) # abs: potential symmetric around wall 

43 

44 for k in range(D): 

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

46 cuda.atomic.add(scalars, (particle, w_id), dist ** 2 * s) # Virial 

47 cuda.atomic.add(scalars, (particle, u_id), u) # Potential enerrgy 

48 lap = numba.float32(1 - D) * s + umm # Laplacian 

49 cuda.atomic.add(scalars, (particle, lap_id), lap) 

50 

51 return 

52 

53 return planar_calculator 

54 

55 

56def setup_planar_interactions(configuration, potential_function, potential_params_list, particles_list, point_list, 

57 normal_vector_list, compute_plan, verbose=True) -> dict: 

58 """ Returns a dictionary with planar interactions """ 

59 D = configuration.D 

60 num_types = len(potential_params_list) 

61 assert len(particles_list) == num_types 

62 assert len(point_list) == num_types 

63 assert len(normal_vector_list) == num_types 

64 

65 total_number_indices = 0 

66 for particles in particles_list: 

67 total_number_indices += particles.shape[0] 

68 

69 if verbose: 

70 print( 

71 f'Setting up planar interactions: {num_types} types, {total_number_indices} particle-plane interactions in total.') 

72 

73 indices = np.zeros((total_number_indices, 2), dtype=np.int32) 

74 params = np.zeros((num_types, 2 * D + len(potential_params_list[0])), dtype=np.float32) 

75 

76 start_index = 0 

77 for interaction_type in range(num_types): 

78 next_start_index = start_index + len(particles_list[interaction_type]) 

79 indices[start_index:next_start_index, 0] = particles_list[interaction_type] 

80 indices[start_index:next_start_index, 1] = interaction_type 

81 start_index = next_start_index 

82 

83 params[interaction_type, 0:D] = point_list[interaction_type] 

84 params[interaction_type, D:2 * D] = normal_vector_list[interaction_type] # Normalize it! 

85 params[interaction_type, 2 * D:] = potential_params_list[interaction_type] 

86 

87 calculator = make_planar_calculator(configuration, potential_function) 

88 interactions = make_fixed_interactions(configuration, calculator, compute_plan, verbose=False) 

89 d_indices = cuda.to_device(indices) 

90 d_params = cuda.to_device(params) 

91 interaction_params = (d_indices, d_params) 

92 

93 return {'interactions': interactions, 'interaction_params': interaction_params}