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

1# fmt: off 

2 

3import numpy as np 

4 

5from ase.geometry import find_mic 

6 

7 

8def rotation_matrix_from_points(m0, m1): 

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

10 RMSD between two set of points. 

11 

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

13 coordinates as columns:: 

14 

15 (x1 x2 x3 ... xN 

16 y1 y2 y3 ... yN 

17 z1 z2 z3 ... zN) 

18 

19 The centeroids should be set to origin prior to 

20 computing the rotation matrix. 

21 

22 The rotation matrix is computed using quaternion 

23 algebra as detailed in:: 

24 

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

26 """ 

27 

28 v0 = np.copy(m0) 

29 v1 = np.copy(m1) 

30 

31 # compute the rotation quaternion 

32 

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

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

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

36 

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

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

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

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

41 

42 F = np.array(f) 

43 

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

45 # eigenvector corresponding to the most 

46 # positive eigenvalue 

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

48 

49 # Rotation matrix from the quaternion q 

50 

51 R = quaternion_to_matrix(q) 

52 

53 return R 

54 

55 

56def quaternion_to_matrix(q): 

57 """Returns a rotation matrix. 

58 

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

60 """ 

61 

62 q0, q1, q2, q3 = q 

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

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

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

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

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

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

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

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

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

72 return np.array(R_q) 

73 

74 

75def minimize_rotation_and_translation(target, atoms): 

76 """Minimize RMSD between atoms and target. 

77 

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

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

80 

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

82 """ 

83 

84 p = atoms.get_positions() 

85 p0 = target.get_positions() 

86 

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

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

89 # account for them? 

90 

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

92 dp_min, _dp_len = find_mic(p - p0, cell=target.cell, pbc=target.pbc) 

93 

94 # add displacement without net translation 

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

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

97 

98 # centeroids to origin 

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

100 p -= c 

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

102 p0 -= c0 

103 

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

105 # Compute rotation matrix 

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

107 

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