1from numpy import einsum, sqrt, eye, copy, amax, sign, prod, linalg, zeros, arange, zeros_like 

2from mango.constants import c 

3 

4 

5class CoM_motion_rm(): 

6 

7 def __init__(self, mass, no_mol): 

8 self.mass = mass 

9 self.no_mol = no_mol 

10 self.stats = [] 

11 

12 def suscep_align(self, sus_data): 

13 pos = sus_data['pos'] 

14 mom = sus_data['mom'] 

15 mag = sus_data['mag'] 

16 

17 rot, inertia, angular = self._rm(pos, mom, mag) 

18 

19 sus_data['rotation'] = rot 

20 sus_data['inertia'] = inertia 

21 sus_data['angular'] = angular 

22 

23 return sus_data 

24 

25 def align(self, xyz_data, sus_data): 

26 if xyz_data is None and sus_data is not None: 

27 self.suscep_align(sus_data) 

28 elif xyz_data is not None: 

29 rot, inertia, angular = self.xyz_align(xyz_data) 

30 

31 if sus_data is not None: 

32 sus_data['pos'][:] = xyz_data[..., 0:3] 

33 sus_data['mom'][:] = xyz_data[..., 3:6] 

34 sus_data['mag'][:] = xyz_data[..., 6:9] 

35 sus_data['rotation'] = rot 

36 sus_data['inertia'] = inertia 

37 sus_data['angular'] = angular 

38 

39 return xyz_data, sus_data 

40 

41 def xyz_align(self, xyz_data): 

42 pos = xyz_data[..., 0:3] 

43 mom = xyz_data[..., 3:6] 

44 mag = xyz_data[..., 6:9] 

45 

46 rot, inertia, angular = self._rm(pos, mom, mag) 

47 

48 return rot, inertia, angular 

49 

50 def inertia_angular(self, xyz_data, sus_data): 

51 pos = sus_data['pos'] 

52 mom = sus_data['mom'] 

53 pos[:], mom[:], pos_com, mom_com = remove_CoMm(pos, mom, self.mass, self.no_mol) 

54 return get_angular(pos, pos_com, mom, mom_com), get_inertia(pos, pos_com, self.mass) 

55 

56 def _rm(self, pos, mom, mag): 

57 

58 pos[:], mom[:], pos_com, mom_com = remove_CoMm(pos, mom, self.mass, self.no_mol) 

59 

60 inertia = get_inertia(pos, pos_com, self.mass) 

61 

62 rot = zeros_like(inertia) 

63 

64 # Kabsch's algorithm 

65 if self.no_mol > 1: 

66 # Alignment of the initial frame (diagonalise the tensor of inertia) 

67 wmat, vmat = linalg.eigh(inertia[:, 0]) 

68 

69 fmat = einsum('aij -> aji', vmat) 

70 

71 print("# Frame alignment (Kabsch)") 

72 

73 target = copy(pos[:, 0]) # Initial configuration 

74 

75 rmsd_max = amax(einsum('aijk,aijk->ai', pos - target[:, None], pos - target[:, None]) / self.no_mol) 

76 

77 print("# Maximum RMSD before alignment: %12.5g" % (rmsd_max)) 

78 

79 rmat = einsum('abxi,axj->abij', pos, target) 

80 vmat, sigma, wmat = linalg.svd(rmat) 

81 umat = zeros((pos.shape[0], pos.shape[1], 3, 3)) 

82 

83 loc = arange(3) 

84 umat[:, :, loc, loc] = 1 

85 umat[:, :, 2, 2] = sign(prod(sigma, axis=-1)) 

86 

87 rmat = einsum('aij, abkj, abck, ablk -> abil', fmat, wmat, umat, vmat) 

88 

89 # Rotation 

90 pos[:] = einsum('abij,abxj->abxi', rmat, pos) 

91 mom[:] = einsum('abij,abxj->abxi', rmat, mom) 

92 mag[:] = einsum('abij,abxj->abxi', rmat, mag) 

93 rot[:] = einsum('abji', rmat) 

94 

95 target[:] = einsum('axk,ajk->ajx', fmat, target) 

96 

97 rmsd_max = amax(einsum('aijk,aijk->ai', pos - target[:, None, :, :], pos - target[:, None, :, :]) / self.no_mol) 

98 

99 print("# Maximum RMSD after alignment: %12.5g" % (rmsd_max)) 

100 

101 pos[:], mom[:], pos_com, mom_com = remove_CoMm(pos, mom, self.mass, self.no_mol) 

102 

103 inertia = get_inertia(pos, pos_com, self.mass) 

104 

105 angular = get_angular(pos, pos_com, mom, mom_com) 

106 

107 return rot, inertia, angular 

108 

109 

110def get_inertia(pos, pos_com, mass): 

111 # Tensor of inertia 

112 mass_scaled = einsum('aijk,j->aijk', pos - pos_com[:, :, None, :], sqrt(mass)) # mass-scaled coordinates 

113 tdot = einsum('aixj,aixk->aijk', mass_scaled, mass_scaled) 

114 

115 return einsum('aixx,jk->aijk', tdot, eye(3)) - tdot 

116 

117 

118def get_angular(pos, pos_com, mom, mom_com): 

119 return einsum('xyz,aijx, aijy ->aijz', c.eijk, pos - pos_com[:, :, None, :], mom - mom_com[:, :, None, :], optimize=True) 

120 

121 

122def remove_CoMm(pos, mom, mass, no_mol): 

123 # Remove Centre of Mass Motion 

124 pos_com = einsum("aijk,j->aik", pos, mass) / einsum("j->", mass) 

125 pos -= pos_com[:, :, None, :] 

126 

127 mom_com = einsum("aijk->aik", mom) / no_mol 

128 mom -= mom_com[:, :, None, :] 

129 

130 return pos, mom, pos_com, mom_com