Coverage for /builds/kinetik161/ase/ase/build/rotate.py: 100.00%

32 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-12-10 11:04 +0000

1import numpy as np 

2from ase.geometry import find_mic 

3 

4 

5def rotation_matrix_from_points(m0, m1): 

6 """Returns a rigid transformation/rotation matrix that minimizes the 

7 RMSD between two set of points. 

8 

9 m0 and m1 should be (3, npoints) numpy arrays with 

10 coordinates as columns:: 

11 

12 (x1 x2 x3 ... xN 

13 y1 y2 y3 ... yN 

14 z1 z2 z3 ... zN) 

15 

16 The centeroids should be set to origin prior to 

17 computing the rotation matrix. 

18 

19 The rotation matrix is computed using quaternion 

20 algebra as detailed in:: 

21 

22 Melander et al. J. Chem. Theory Comput., 2015, 11,1055 

23 """ 

24 

25 v0 = np.copy(m0) 

26 v1 = np.copy(m1) 

27 

28 # compute the rotation quaternion 

29 

30 R11, R22, R33 = np.sum(v0 * v1, axis=1) 

31 R12, R23, R31 = np.sum(v0 * np.roll(v1, -1, axis=0), axis=1) 

32 R13, R21, R32 = np.sum(v0 * np.roll(v1, -2, axis=0), axis=1) 

33 

34 f = [[R11 + R22 + R33, R23 - R32, R31 - R13, R12 - R21], 

35 [R23 - R32, R11 - R22 - R33, R12 + R21, R13 + R31], 

36 [R31 - R13, R12 + R21, -R11 + R22 - R33, R23 + R32], 

37 [R12 - R21, R13 + R31, R23 + R32, -R11 - R22 + R33]] 

38 

39 F = np.array(f) 

40 

41 w, V = np.linalg.eigh(F) 

42 # eigenvector corresponding to the most 

43 # positive eigenvalue 

44 q = V[:, np.argmax(w)] 

45 

46 # Rotation matrix from the quaternion q 

47 

48 R = quaternion_to_matrix(q) 

49 

50 return R 

51 

52 

53def quaternion_to_matrix(q): 

54 """Returns a rotation matrix. 

55 

56 Computed from a unit quaternion Input as (4,) numpy array. 

57 """ 

58 

59 q0, q1, q2, q3 = q 

60 R_q = [[q0**2 + q1**2 - q2**2 - q3**2, 

61 2 * (q1 * q2 - q0 * q3), 

62 2 * (q1 * q3 + q0 * q2)], 

63 [2 * (q1 * q2 + q0 * q3), 

64 q0**2 - q1**2 + q2**2 - q3**2, 

65 2 * (q2 * q3 - q0 * q1)], 

66 [2 * (q1 * q3 - q0 * q2), 

67 2 * (q2 * q3 + q0 * q1), 

68 q0**2 - q1**2 - q2**2 + q3**2]] 

69 return np.array(R_q) 

70 

71 

72def minimize_rotation_and_translation(target, atoms): 

73 """Minimize RMSD between atoms and target. 

74 

75 Rotate and translate atoms to best match target. Disregards rotation if PBC 

76 are found. Does not accound for changes in the cell. For more details, see:: 

77 

78 Melander et al. J. Chem. Theory Comput., 2015, 11,1055 

79 """ 

80 

81 p = atoms.get_positions() 

82 p0 = target.get_positions() 

83 

84 if sum(atoms.pbc) != 0: 

85 # maybe we can raise a warning about cell changes here since we don't 

86 # account for them? 

87 

88 # is this the best form of *find_mic version to use? 

89 dp_min, dp_len = find_mic(p - p0, cell=target.cell, pbc=target.pbc) 

90 

91 # add displacement without net translation 

92 p = p0 + dp_min - np.mean(dp_min, axis=0) 

93 R = np.eye(3) # null rotation 

94 

95 # centeroids to origin 

96 c = np.mean(p, axis=0) 

97 p -= c 

98 c0 = np.mean(p0, axis=0) 

99 p0 -= c0 

100 

101 if sum(atoms.pbc) == 0: 

102 # Compute rotation matrix 

103 R = rotation_matrix_from_points(p.T, p0.T) 

104 

105 atoms.set_positions(np.dot(p, R.T) + c0)