Coverage for gamdpy/interactions/gravity.py: 74%

1 

2import numpy as np 

3from numba import cuda 

4import numba 

5import math 

6from .make_fixed_interactions import make_fixed_interactions 

7 

8# Abstract Base Class and type annotation 

9from .interaction import Interaction 

10from gamdpy import Configuration 

11 

12class Gravity(Interaction): 

13 """Adding a gravitational-like force on particles. 

14 

15 Parameters 

16 ---------- 

17 (a) A list of particle indices on which the forces act and associated list of forces  

18 (b) A list of particle types one which the forces act and associated list of forces 

19 

20 At the moment the force will act in the x-direction  

21 

22 See examples/poiseuille.py 

23 """ 

24 

25 def __init__(self): 

26 self.indices_set = False 

27 

28 

29 def set_gravity_from_lists(self, particle_indices, forces): 

30 

31 if len(forces) != len(particle_indices): 

32 raise ValueError("Force and particle index arrays must have same length") 

33 

34 force_array, indices_array = [], [] 

35 

36 for n in range(len(particle_indices)): 

37 indices_array.append( [n, pindices[n]] ) 

38 force_array.append( forces[n] ) 

39 

40 self.force_array = np.array(force_array, dtype=np.float32) 

41 self.indices_array = np.array(indices_array, dtype=np.int32) 

42 

43 self.indices_set = True 

44 

45 

46 def set_gravity_from_types(self, particle_types, forces, configuration): 

47 

48 ntypes, nforces = len(particle_types), len(forces) 

49 

50 if ntypes != nforces: 

51 raise ValueError("Force and particle type arrays must have same length") 

52 

53 force_array, indices_array = [], [] 

54 

55 counter = 0 

56 for n in range(configuration.N): 

57 for m in range(ntypes): 

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

59 indices_array.append( [counter, n] ) 

60 force_array.append( forces[m] ) 

61 counter = counter + 1 

62 break 

63 

64 self.force_array = np.array(force_array, dtype=np.float32) 

65 self.indices_array = np.array(indices_array, dtype=np.int32) 

66 

67 self.indices_set = True 

68 

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

70 if self.indices_set == False: 

71 raise ValueError("Indices not defined") 

72 

73 self.d_pindices = cuda.to_device(self.indices_array) 

74 self.d_force = cuda.to_device(self.force_array); 

75 

76 return (self.d_pindices, self.d_force) 

77 

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

79 # Unpack parameters from configuration and compute_plan 

80 D, N = configuration.D, configuration.N 

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

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

83 

84 compute_u = compute_flags['U'] # PE should be included ?!?! 

85 # Note virial, lapacian, stresses are zero for gravity 

86 f_id = configuration.vectors.indices['f'] 

87 

88 def gravity_calculator(vectors, scalars, ptype, sim_box, indices, values): 

89 

90 f = vectors[f_id][indices[1]] 

91 

92 f[0] = f[0] + values[indices[0]] 

93 

94 return 

95 

96 return make_fixed_interactions(configuration, gravity_calculator, compute_plan, verbose=False) 

97 

98 

99