Coverage for src / sdynpy / signal_processing / sdynpy_rotation.py: 27%

1# -*- coding: utf-8 -*- 

2""" 

3Functions for dealing with the geometry of rotations. 

4""" 

5""" 

6Copyright 2022 National Technology & Engineering Solutions of Sandia, 

7LLC (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S. 

8Government retains certain rights in this software. 

9 

10This program is free software: you can redistribute it and/or modify 

11it under the terms of the GNU General Public License as published by 

12the Free Software Foundation, either version 3 of the License, or 

13(at your option) any later version. 

14 

15This program is distributed in the hope that it will be useful, 

16but WITHOUT ANY WARRANTY; without even the implied warranty of 

17MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

18GNU General Public License for more details. 

19 

20You should have received a copy of the GNU General Public License 

21along with this program. If not, see <https://www.gnu.org/licenses/>. 

22""" 

23 

24import numpy as np 

25from scipy.optimize import NonlinearConstraint 

26 

27 

28def cross_mat(v): 

29 """ 

30 Return matrix representation of the cross product of the given vector 

31 

32 Parameters 

33 ---------- 

34 v : np.ndarray 

35 3-vector that will be assembled into the matrix. 

36 

37 Returns 

38 ------- 

39 cross_matrix : np.ndarray 

40 Matrix representing the cross product. np.cross(a,b) is equivalent to 

41 cross_mat(a)@b 

42 

43 """ 

44 v = np.array(v).flatten() 

45 return np.array([[0, -v[2], v[1]], 

46 [v[2], 0, -v[0]], 

47 [-v[1], v[0], 0]]) 

48 

49 

50def R(axis, angle, degrees=False): 

51 """ 

52 Create a rotation matrix consisting of a rotation about an axis 

53 

54 Parameters 

55 ---------- 

56 axis : index or np.ndarray 

57 The axis of rotation. Specifying 0, 1, or 2 will construct a rotation 

58 about the X, Y, or Z axes, respectively. Alternatively, the axis of 

59 rotation can be specficied as a 3-vector. 

60 angle : float 

61 The angle of rotation 

62 degrees : bool, optional 

63 True if the angle value is specified in degrees, false if radians. 

64 The default is False. 

65 

66 Raises 

67 ------ 

68 ValueError 

69 If an improper axis is specified. 

70 

71 Returns 

72 ------- 

73 rotmat : np.ndarray 

74 A 3x3 rotation matrix. 

75 

76 """ 

77 rotmat = None 

78 if degrees: 

79 angle *= np.pi / 180 

80 s = np.sin(angle) 

81 c = np.cos(angle) 

82 try: 

83 if axis == 0: 

84 rotmat = np.array([[1, 0, 0], 

85 [0, c, -s], 

86 [0, s, c]]) 

87 elif axis == 1: 

88 rotmat = np.array([[c, 0, s], 

89 [0, 1, 0], 

90 [-s, 0, c]]) 

91 elif axis == 2: 

92 rotmat = np.array([[c, -s, 0], 

93 [s, c, 0], 

94 [0, 0, 1]]) 

95 else: 

96 axis = np.reshape(axis, (3, 1)) / np.linalg.norm(axis) 

97 rotmat = np.eye(3) * c + s * cross_mat(axis) + (1 - c) * (axis @ axis.T) 

98 except ValueError: 

99 axis = np.reshape(axis, (3, 1)) / np.linalg.norm(axis) 

100 rotmat = np.eye(3) * c + s * cross_mat(axis) + (1 - c) * (axis @ axis.T) 

101 if rotmat is None: 

102 raise ValueError('Invalid Axis {:}'.format(axis)) 

103 return rotmat 

104 

105 

106def quaternion_to_rotation_matrix(quat): 

107 output = np.zeros(quat.shape[:-1] + (3, 3)) 

108 q0 = quat[..., 0] 

109 q1 = quat[..., 1] 

110 q2 = quat[..., 2] 

111 q3 = quat[..., 3] 

112 

113 # From https://automaticaddison.com/how-to-convert-a-quaternion-to-a-rotation-matrix/ 

114 # First row of the rotation matrix 

115 output[..., 0, 0] = 2 * (q0 * q0 + q1 * q1) - 1 

116 output[..., 0, 1] = 2 * (q1 * q2 - q0 * q3) 

117 output[..., 0, 2] = 2 * (q1 * q3 + q0 * q2) 

118 # Second row of the rotation matrix 

119 output[..., 1, 0] = 2 * (q1 * q2 + q0 * q3) 

120 output[..., 1, 1] = 2 * (q0 * q0 + q2 * q2) - 1 

121 output[..., 1, 2] = 2 * (q2 * q3 - q0 * q1) 

122 # Third row of the rotation matrix 

123 output[..., 2, 0] = 2 * (q1 * q3 - q0 * q2) 

124 output[..., 2, 1] = 2 * (q2 * q3 + q0 * q1) 

125 output[..., 2, 2] = 2 * (q0 * q0 + q3 * q3) - 1 

126 

127 return output 

128 

129 

130unit_magnitude_constraint = NonlinearConstraint(np.linalg.norm, 1, 1) 

131 

132 

133def lstsq_rigid_transform(x, y, w=None): 

134 """ Computes the Transform between two point sets such that 

135 y = Rx + t 

136 

137 This is a least squares methods as descibed in 

138 "Least-Squares Rigid Motion Using SVD", O. Sorkine-Hornung, M. 

139 Rabinovich, Department of Computer Science, ETH Zurich, Jan 16, 2017 

140 Note found online at https://igl.ethz.ch/projects/ARAP/svd_rot.pdf 

141 

142 INPUT 

143 x,y = the two point sets as [x,y,z] column vectors 

144 sizes are (3,n) for n points 

145 w = weighting vector for each point, wi>0 

146 default is uniform weighting of wi=1 

147 size is (1,n) for n points 

148 

149 OUTPUT 

150 R,t = [(3,3),(3,1)] transformation parameters such that 

151 y = Rx + t 

152 """ 

153 if w is None: 

154 w = np.ones((x.shape[-1], 1)).T 

155 # default weighting is uniform w=1 

156 

157 W = np.diag(w[0]) 

158 

159 # Find the weighted centroids (for w=1, this is the mean) 

160 xbar = np.sum(w * x, axis=-1) / np.sum(w, axis=1) 

161 ybar = np.sum(w * y, axis=-1) / np.sum(w, axis=1) 

162 

163 # Center the points 

164 X = x - xbar[..., np.newaxis] 

165 Y = y - ybar[..., np.newaxis] 

166 

167 # Calculate the Covariance 

168 Cov = X @ W @ np.moveaxis(Y, -1, -2) 

169 

170 # Take the SVD of the Covariance matrix 

171 U, S, VH = np.linalg.svd(Cov) 

172 V = np.moveaxis(VH, -1, -2) # numpy's SVD gives you back V', not V like Matlab 

173 det = np.linalg.det(V @ np.moveaxis(U, -1, -2)) 

174 D = np.broadcast_to(np.eye(3), det.shape+(3, 3)).copy() 

175 D[..., -1, -1] = det 

176 

177 R = V @ D @ np.moveaxis(U, -1, -2) 

178 t = ybar[..., np.newaxis] - R @ xbar[..., np.newaxis] 

179 

180 return R, t 

181 

182 

183def cross_mat_vectorized(rvec): 

184 zero = np.zeros(rvec.shape[:-1]) 

185 return np.moveaxis(np.moveaxis( 

186 np.array([ 

187 [zero, -rvec[..., 2], rvec[..., 1]], 

188 [rvec[..., 2], zero, -rvec[..., 0]], 

189 [-rvec[..., 1], rvec[..., 0], zero] 

190 ]), 0, -1), 0, -1) 

191 

192 

193def rodrigues_to_matrix(rvec, threshold=0.000001): 

194 theta = np.linalg.norm(rvec, axis=-1, keepdims=True) 

195 rvec = rvec / theta 

196 R = ( 

197 np.cos(theta[..., np.newaxis]) * np.eye(3) 

198 + (1 - np.cos(theta[..., np.newaxis])) * rvec[..., np.newaxis] @ rvec[..., np.newaxis, :] 

199 + np.sin(theta[..., np.newaxis]) * cross_mat_vectorized(rvec)) 

200 R[theta[..., 0] < threshold] = np.eye(3) 

201 return R 

202 

203 

204def matrix_to_rodrigues(R, threshold=0.000001): 

205 A = (R - np.moveaxis(R, -2, -1)) / 2 

206 rho = np.moveaxis(np.array((A[..., 2, 1], A[..., 0, 2], A[..., 1, 0])), 0, -1) 

207 s = np.linalg.norm(rho, axis=-1) 

208 c = (R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] - 1) / 2 

209 u = np.zeros(rho.shape) 

210 u[s >= threshold] = rho[s >= threshold] / s[s >= threshold][:, np.newaxis] 

211 u[(s < threshold) & (c > 1 - threshold)] = 0 

212 theta = np.arctan2(s, c) 

213 rvec = u * theta[..., np.newaxis] 

214 # Now handle the cases where s = 0 and c = -1 

215 special_cases = (s < threshold) & (c < -1 + threshold) 

216 for index in zip(*np.where(special_cases)): 

217 thisR = R[index] 

218 RpI = thisR + np.eye(3) 

219 best_column = np.linalg.norm(RpI, axis=0).argmax() 

220 v = RpI[:, best_column] 

221 thisu = v / np.linalg.norm(v) 

222 if ((abs(thisu[0]) < threshold and abs(thisu[1]) < threshold and thisu[2] < 0) or 

223 (abs(thisu[0] < threshold and thisu[1] < 0)) or 

224 (thisu[0] < 0)): 

225 rvec[index] = -thisu * np.pi 

226 else: 

227 rvec[index] = thisu * np.pi 

228 return rvec