Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_matrix_from_tait_bryan_angles_extrinsic_xyz(self):
tait_bryan_angles = gs.array([0., 0., 0.])
result = self.group.matrix_from_tait_bryan_angles_extrinsic_xyz(
tait_bryan_angles)
expected = gs.eye(3)
self.assertAllClose(result, expected)
angle = gs.pi / 6.
cos_angle = gs.cos(angle)
sin_angle = gs.sin(angle)
tait_bryan_angles = gs.array([angle, 0., 0.])
result = self.group.matrix_from_tait_bryan_angles_extrinsic_xyz(
tait_bryan_angles)
expected = gs.array([[cos_angle, - sin_angle, 0.],
[sin_angle, cos_angle, 0.],
[0., 0., 1.]])
self.assertAllClose(result, expected)
tait_bryan_angles = gs.array([0., angle, 0.])
result = self.group.matrix_from_tait_bryan_angles_extrinsic_xyz(
tait_bryan_angles)
expected = gs.array([[cos_angle, 0., sin_angle],
[0., 1., 0.],
[- sin_angle, 0., cos_angle]])
gs_array = gs.zeros((2, 3, 3))
np_array[:, 0, 1] = 44.
gs_array = gs.assignment(
gs_array, 44., (0, 1), axis=0)
self.assertAllCloseToNp(gs_array, np_array)
n_samples = 3
theta = _np.random.rand(5)
phi = _np.random.rand(5)
np_array = _np.zeros((n_samples, 5, 4))
gs_array = gs.array(np_array)
np_array[0, :, 0] = gs.cos(theta) * gs.cos(phi)
np_array[0, :, 1] = - gs.sin(theta) * gs.sin(phi)
gs_array = gs.assignment(
gs_array, gs.cos(theta) * gs.cos(phi), (0, 0), axis=1)
gs_array = gs.assignment(
gs_array, - gs.sin(theta) * gs.sin(phi), (0, 1), axis=1)
self.assertAllCloseToNp(gs_array, np_array)
"""
n_rot_mats, _, _ = rot_mat.shape
trace = gs.trace(rot_mat, axis1=1, axis2=2)
trace = gs.to_ndarray(trace, to_ndim=2, axis=1)
trace_num = gs.clip(trace, -1, 3)
angle = gs.arccos(0.5 * (trace_num - 1))
rot_mat_transpose = gs.transpose(rot_mat, axes=(0, 2, 1))
rot_vec_not_pi = self.vector_from_skew_matrix(
rot_mat - rot_mat_transpose)
mask_0 = gs.cast(gs.isclose(angle, 0.), gs.float32)
mask_pi = gs.cast(gs.isclose(angle, gs.pi, atol=1e-2), gs.float32)
mask_else = (1 - mask_0) * (1 - mask_pi)
numerator = 0.5 * mask_0 + angle * mask_else
denominator = (1 - angle ** 2 / 6) * mask_0 + 2 * gs.sin(
angle) * mask_else + mask_pi
rot_vec_not_pi = rot_vec_not_pi * numerator / denominator
vector_outer = 0.5 * (gs.eye(3) + rot_mat)
gs.set_diag(
vector_outer, gs.maximum(
0., gs.diagonal(vector_outer, axis1=1, axis2=2)))
squared_diag_comp = gs.diagonal(vector_outer, axis1=1, axis2=2)
diag_comp = gs.sqrt(squared_diag_comp)
norm_line = gs.linalg.norm(vector_outer, axis=2)
max_line_index = gs.argmax(norm_line, axis=1)
selected_line = gs.get_slice(
vector_outer, (range(n_rot_mats), max_line_index))
signs = gs.sign(selected_line)
rot_vec_pi = angle * signs * diag_comp
n_rot_vecs, _ = rot_vec.shape
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
rotation_axis = gs.zeros_like(rot_vec)
mask_0 = gs.isclose(angle, 0.)
mask_0 = gs.squeeze(mask_0, axis=1)
mask_not_0 = ~mask_0
rotation_axis[mask_not_0] = rot_vec[mask_not_0] / angle[mask_not_0]
n_quaternions, _ = rot_vec.shape
quaternion = gs.zeros((n_quaternions, 4))
quaternion[:, :1] = gs.cos(angle / 2)
quaternion[:, 1:] = gs.sin(angle / 2) * rotation_axis[:]
return quaternion
tait_bryan_angles : array-like, shape=[..., 3]
Returns
-------
rot_mat : array-like, shape=[..., n, n]
"""
n_tait_bryan_angles, _ = tait_bryan_angles.shape
rot_mat = []
angle_1 = tait_bryan_angles[:, 0]
angle_2 = tait_bryan_angles[:, 1]
angle_3 = tait_bryan_angles[:, 2]
for i in range(n_tait_bryan_angles):
cos_angle_1 = gs.cos(angle_1[i])
sin_angle_1 = gs.sin(angle_1[i])
cos_angle_2 = gs.cos(angle_2[i])
sin_angle_2 = gs.sin(angle_2[i])
cos_angle_3 = gs.cos(angle_3[i])
sin_angle_3 = gs.sin(angle_3[i])
column_1 = [[cos_angle_2 * cos_angle_3],
[(cos_angle_1 * sin_angle_3
+ cos_angle_3 * sin_angle_1 * sin_angle_2)],
[(sin_angle_1 * sin_angle_3
- cos_angle_1 * cos_angle_3 * sin_angle_2)]]
column_2 = [[- cos_angle_2 * sin_angle_3],
[(cos_angle_1 * cos_angle_3
- sin_angle_1 * sin_angle_2 * sin_angle_3)],
[(cos_angle_3 * sin_angle_1
+ cos_angle_1 * sin_angle_2 * sin_angle_3)]]
"""
n_rot_mats, _, _ = rot_mat.shape
trace = gs.trace(rot_mat, axis1=1, axis2=2)
trace = gs.to_ndarray(trace, to_ndim=2, axis=1)
trace_num = gs.clip(trace, -1, 3)
angle = gs.arccos(0.5 * (trace_num - 1))
rot_mat_transpose = gs.transpose(rot_mat, axes=(0, 2, 1))
rot_vec_not_pi = self.vector_from_skew_matrix(
rot_mat - rot_mat_transpose)
mask_0 = gs.cast(gs.isclose(angle, 0.), gs.float32)
mask_pi = gs.cast(gs.isclose(angle, gs.pi, atol=1e-2), gs.float32)
mask_else = (1 - mask_0) * (1 - mask_pi)
numerator = 0.5 * mask_0 + angle * mask_else
denominator = (1 - angle ** 2 / 6) * mask_0 + 2 * gs.sin(
angle) * mask_else + mask_pi
rot_vec_not_pi = rot_vec_not_pi * numerator / denominator
vector_outer = 0.5 * (gs.eye(3) + rot_mat)
gs.set_diag(
vector_outer, gs.maximum(
0., gs.diagonal(vector_outer, axis1=1, axis2=2)))
squared_diag_comp = gs.diagonal(vector_outer, axis1=1, axis2=2)
diag_comp = gs.sqrt(squared_diag_comp)
norm_line = gs.linalg.norm(vector_outer, axis=2)
max_line_index = gs.argmax(norm_line, axis=1)
selected_line = gs.get_slice(
vector_outer, (range(n_rot_mats), max_line_index))
signs = gs.sign(selected_line)
rot_vec_pi = angle * signs * diag_comp
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
mask_0 = gs.isclose(angle, 0.)
mask_not_0 = ~mask_0
rotation_axis = gs.divide(
rot_vec,
angle
* gs.cast(mask_not_0, gs.float32)
+ gs.cast(mask_0, gs.float32))
quaternion = gs.concatenate(
(gs.cos(angle / 2),
gs.sin(angle / 2) * rotation_axis[:]),
axis=1)
return quaternion
Returns
-------
rot_mat : array-like, shape=[n_samples, n, n]
"""
n_tait_bryan_angles, _ = tait_bryan_angles.shape
rot_mat = gs.zeros((n_tait_bryan_angles,) + (self.n,) * 2)
angle_1 = tait_bryan_angles[:, 0]
angle_2 = tait_bryan_angles[:, 1]
angle_3 = tait_bryan_angles[:, 2]
for i in range(n_tait_bryan_angles):
cos_angle_1 = gs.cos(angle_1[i])
sin_angle_1 = gs.sin(angle_1[i])
cos_angle_2 = gs.cos(angle_2[i])
sin_angle_2 = gs.sin(angle_2[i])
cos_angle_3 = gs.cos(angle_3[i])
sin_angle_3 = gs.sin(angle_3[i])
column_1 = [[cos_angle_2 * cos_angle_3],
[(cos_angle_1 * sin_angle_3
+ cos_angle_3 * sin_angle_1 * sin_angle_2)],
[(sin_angle_1 * sin_angle_3
- cos_angle_1 * cos_angle_3 * sin_angle_2)]]
column_2 = [[- cos_angle_2 * sin_angle_3],
[(cos_angle_1 * cos_angle_3
- sin_angle_1 * sin_angle_2 * sin_angle_3)],
[(cos_angle_3 * sin_angle_1
+ cos_angle_1 * sin_angle_2 * sin_angle_3)]]
column_3 = [[sin_angle_2],
angle = gs.linalg.norm(rot_vec, axis=1)
angle = gs.to_ndarray(angle, to_ndim=2, axis=1)
mask_0 = gs.isclose(angle, 0.)
mask_not_0 = ~mask_0
rotation_axis = gs.divide(
rot_vec,
angle
* gs.cast(mask_not_0, gs.float32)
+ gs.cast(mask_0, gs.float32))
quaternion = gs.concatenate(
(gs.cos(angle / 2),
gs.sin(angle / 2) * rotation_axis[:]),
axis=1)
return quaternion