Coverage for gamdpy/simulation_boxes/orthorhombic.py: 99%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4Created on Thu Jun 13 11:41:24 2024 

5 

6@author: nbailey 

7""" 

8 

9import numpy as np 

10import numba 

11from numba import cuda 

12from .simulationbox import SimulationBox 

13 

14class Orthorhombic(SimulationBox): 

15 """ Standard rectangular simulation box class  

16 

17 Example 

18 ------- 

19 

20 >>> import gamdpy as gp 

21 >>> import numpy as np 

22 >>> simbox = gp.Orthorhombic(D=3, lengths=np.array([3,4,5])) 

23 

24 """ 

25 def __init__(self, D, lengths): 

26 self.D = D 

27 self.data_array = np.array(lengths, dtype=np.float32) # ensure single precision 

28 self.len_sim_box_data = D # not true for other Simbox classes. Want to remove this and just use len(self.data_array) 

29 return 

30 

31 def get_name(self): 

32 return "Orthorhombic" 

33 

34 

35 def copy_to_device(self): 

36 self.d_data = cuda.to_device(self.data_array) 

37 

38 def copy_to_host(self): 

39 self.data_array = self.d_data.copy_to_host() 

40 

41 def get_dist_sq_dr_function(self): 

42 """Generates function dist_sq_dr which computes displacement and distance squared for one neighbor """ 

43 D = self.D 

44 def dist_sq_dr_function(ri, rj, sim_box, dr): 

45 

46 ''' Returns the squared distance between ri and rj applying MIC and saves ri-rj in dr ''' 

47 dist_sq = numba.float32(0.0) 

48 for k in range(D): 

49 dr[k] = ri[k] - rj[k] 

50 box_k = sim_box[k] 

51 dr[k] += (-box_k if numba.float32(2.0) * dr[k] > +box_k else 

52 (+box_k if numba.float32(2.0) * dr[k] < -box_k else numba.float32(0.0))) # MIC 

53 dist_sq = dist_sq + dr[k] * dr[k] 

54 return dist_sq 

55 

56 return dist_sq_dr_function 

57 

58 def get_dist_sq_function(self): 

59 """Generates.function dist_sq_function which computes distance squared for one neighbor """ 

60 D = self.D 

61 def dist_sq_function(ri, rj, sim_box): 

62 ''' Returns the squared distance between ri and rj applying MIC''' 

63 dist_sq = numba.float32(0.0) 

64 for k in range(D): 

65 dr_k = ri[k] - rj[k] 

66 box_k = sim_box[k] 

67 dr_k += (-box_k if numba.float32(2.0) * dr_k > +box_k else 

68 (+box_k if numba.float32(2.0) * dr_k < -box_k else numba.float32(0.0))) # MIC 

69 dist_sq = dist_sq + dr_k * dr_k 

70 return dist_sq 

71 

72 return dist_sq_function 

73 

74 def get_apply_PBC(self): 

75 D = self.D 

76 def apply_PBC(r, image, sim_box): 

77 for k in range(D): 

78 if r[k] * numba.float32(2.0) > +sim_box[k]: 

79 r[k] -= sim_box[k] 

80 image[k] += 1 

81 if r[k] * numba.float32(2.0) < -sim_box[k]: 

82 r[k] += sim_box[k] 

83 image[k] -= 1 

84 return 

85 return apply_PBC 

86 

87 def get_lengths(self): 

88 """ Return the box lengths as a numpy array """ 

89 return self.data_array.copy() 

90 

91 def get_volume(self): 

92 """ Return the box volume """ 

93 #self.copy_to_host() # not necessary if volume is fixed and if not fixed then presumably stuff like normalizing stress by volume should be done in the device anyway 

94 return self.get_volume_function()(self.data_array) 

95 

96 def get_volume_function(self): 

97 """ Return the box volume """ 

98 D = self.D 

99 def volume(sim_box): 

100 ''' Returns volume of the rectangular box ''' 

101 vol = sim_box[0] 

102 for i in range(1,D): 

103 vol *= sim_box[i] 

104 return vol 

105 return volume 

106 

107 def scale(self, scale_factor): 

108 """ Scale the box lengths by scale_factor """ 

109 self.data_array *= scale_factor 

110 

111 #def get_dist_moved_sq_function(self): 

112 # D = self.D 

113 # def dist_moved_sq_function(r_current, r_last, sim_box, sim_box_last): 

114 # ''' Returns squared distance between vectors r_current and r_last ''' 

115 # dist_sq = numba.float32(0.0) 

116 # for k in range(D): 

117 # dr_k = r_current[k] - r_last[k] 

118 # box_k = sim_box[k] 

119 # dr_k += (-box_k if numba.float32(2.0) * dr_k > +box_k else 

120 # (+box_k if numba.float32(2.0) * dr_k < -box_k else numba.float32(0.0))) # MIC 

121 # dist_sq = dist_sq + dr_k * dr_k 

122 

123 # return dist_sq 

124 # return dist_moved_sq_function 

125 

126 def get_dist_moved_exceeds_limit_function(self): 

127 D = self.D 

128 

129 def dist_moved_exceeds_limit_function(r_current, r_last, sim_box, sim_box_last, skin, cut): 

130 """ Returns True if squared distance between r_current and r_last exceeds half skin. 

131 Parameters sim_box_last and cut are not used here, but are needed for the Lees-Edwards type of Simbox""" 

132 dist_sq = numba.float32(0.0) 

133 for k in range(D): 

134 dr_k = r_current[k] - r_last[k] 

135 box_k = sim_box[k] 

136 dr_k += (-box_k if numba.float32(2.0) * dr_k > +box_k else 

137 (+box_k if numba.float32(2.0) * dr_k < -box_k else numba.float32(0.0))) # MIC 

138 dist_sq = dist_sq + dr_k * dr_k 

139 

140 return dist_sq > skin*skin*numba.float32(0.25) 

141 return dist_moved_exceeds_limit_function 

142 

143 def get_loop_x_addition(self): 

144 return 0 

145 

146 def get_loop_x_shift_function(self): 

147 

148 def loop_x_shift_function(sim_box, cell_length_x): 

149 return 0 

150 return loop_x_shift_function