1from functools import wraps 

2from numpy import floor, ones_like, einsum, sqrt, absolute, maximum, empty 

3 

4from numpy.core.multiarray import c_einsum 

5 

6from mango.debug import debug 

7from mango.constants import c 

8 

9 

10class periodic(): 

11 """ 

12 Periodic Boundary Conditions. 

13 

14 The positions are only wrapped at the end of a single step 

15 because the values are not used until they are saved to file. 

16 The distances are always wrapped because they are used during a single step. 

17 """ 

18 __slots__ = ['boxsize', 'pbc', 'posones', 'sigma', 'sigma_err', 'tri', 'dist', 

19 'wrap', 'getdist', 'DBG'] 

20 

21 @debug(["distance"]) 

22 def __init__(self): 

23 pass 

24 

25 def setup(self, boxsize, pbc, pos, sigma, tri): 

26 """ 

27 Set up periodic conditions. 

28 

29 Not using __init__ for this allows more decorator flexibility 

30 """ 

31 self.boxsize = boxsize 

32 self.pbc = pbc 

33 self.reset_wrap() 

34 self.posones = ones_like(pos) 

35 self.sigma = sigma 

36 self.sigma_err = self.sigma * 0.7 

37 self.tri = tri 

38 self.dist = empty((pos.shape[-2], pos.shape[-2], 3)) 

39 # self.dist_mat = empty((self.nm, self.nm)) 

40 return self.distance(pos) 

41 

42 def reset_wrap(self): 

43 """Start false.""" 

44 if self.pbc: 

45 self.getdist = self._wrapped_dist 

46 self.wrap = self._wrap2 

47 else: 

48 self.getdist = self._dist 

49 self.wrap = self._nowrap 

50 

51 def setup_full(self, pos): 

52 """ 

53 See setup. 

54 

55 For the full trajectory. 

56 """ 

57 self.getdist = self._dist_full 

58 self.posones = ones_like(pos) 

59 self.dist = empty((*pos.shape[:2], pos.shape[-2], 3)) 

60 # self.dist_mat = empty((self.nm, self.nm)) 

61 return self.distance_full(pos) 

62 

63 def __call__(self, func_to_decorate): 

64 """ 

65 Periodic Boundary condition decorator. 

66 

67 Position Wrapping 

68 

69 * This will provide and adjust the distances and positions of the atoms 

70 based on boundary condition selection 

71 

72 [REF Computer simulation of liquids Allen and Tildesley pg 30] 

73 

74 """ 

75 @wraps(func_to_decorate) 

76 def wrapper(*args, **kw): 

77 result = func_to_decorate(*args, **kw) 

78 

79 args[0].pos[:] = self.wrap(args[0].pos) 

80 args[0].dist[:], args[0].dist_matrix[:] = self.distance(args[0].pos) 

81 

82 return result 

83 return wrapper 

84 

85 @debug(['boundaries']) 

86 def distance(self, pos): 

87 """ 

88 Provide a distance matrix for both periodic and non periodic conditions. 

89 

90 Returns 

91 ------- 

92 dist: array 

93 All pairs distances n(n+1)/2 by 3 

94 dist_mat: array 

95 distance matrix n by n 

96 only ever used as its inverse so calculated as inverse 

97 

98 """ 

99 self.dist[:] = self.getdist(pos) 

100 

101 # if any(absolute(dist_mat[self.tri]) < self.sigma_err): 

102 # Error.count("W", "Particles overlap") 

103 # print(self.sigma, self.sigma_err, absolute(dist_mat[self.tri]) - self.sigma_err);exit(0) 

104 

105 return self.dist, 1 / maximum(sqrt(c_einsum('ijz,ijz->ij', self.dist, self.dist)), c.EPS_val) 

106 

107 def distance_full(self, pos): 

108 """ 

109 See distance. 

110 

111 For the full trajectory 

112 """ 

113 self.dist[:] = self.getdist(pos) 

114 

115 return self.dist, 1 / maximum(sqrt(c_einsum('aijz,aijz->aij', self.dist, self.dist)), c.EPS_val) 

116 

117 def _wrap1(self, coords): 

118 self.wrap = self._wrap2 

119 return self.wrapping(coords) 

120 

121 def _wrap2(self, coords): 

122 self.wrap = self._wrap1 

123 return coords 

124 

125 @staticmethod 

126 def _nowrap(coords): 

127 return coords 

128 

129 def _wrapped_dist(self, pos): 

130 """ 

131 Wrap distance every time. 

132 

133 Positions only need to be wrapped once per loop whereas the wrapped distances are always needed 

134 """ 

135 return self.wrapping(self._dist(pos)) 

136 

137 def _dist(self, pos): 

138 dummy = self.posones 

139 return c_einsum("iz,jz->ijz", pos, dummy) - c_einsum("iz,jz->jiz", pos, dummy) 

140 

141 def _dist_full(self, pos): 

142 dummy = self.posones 

143 return c_einsum("aiz,ajz->aijz", pos, dummy) - c_einsum("aiz,ajz->ajiz", pos, dummy) 

144 

145 def wrapping(self, pos): 

146 """Simplistic wrapping of coordinates.""" 

147 return pos - self.boxsize * floor(pos / self.boxsize + 0.5)