mirror of
https://github.com/3b1b/manim.git
synced 2025-07-28 20:43:56 +08:00
Have rotate_vector return same type as input
This commit is contained in:
@ -63,16 +63,20 @@ def rotate_vector(vector, angle, axis=OUT):
|
|||||||
if len(vector) == 2:
|
if len(vector) == 2:
|
||||||
# Use complex numbers...because why not
|
# Use complex numbers...because why not
|
||||||
z = complex(*vector) * np.exp(complex(0, angle))
|
z = complex(*vector) * np.exp(complex(0, angle))
|
||||||
return np.array([z.real, z.imag])
|
result = [z.real, z.imag]
|
||||||
elif len(vector) == 3:
|
elif len(vector) == 3:
|
||||||
# Use quaternions...because why not
|
# Use quaternions...because why not
|
||||||
quat = quaternion_from_angle_axis(angle, axis)
|
quat = quaternion_from_angle_axis(angle, axis)
|
||||||
quat_inv = quaternion_conjugate(quat)
|
quat_inv = quaternion_conjugate(quat)
|
||||||
product = quaternion_mult(quat, [0, *vector], quat_inv)
|
product = quaternion_mult(quat, [0, *vector], quat_inv)
|
||||||
return product[1:]
|
result = product[1:]
|
||||||
else:
|
else:
|
||||||
raise Exception("vector must be of dimension 2 or 3")
|
raise Exception("vector must be of dimension 2 or 3")
|
||||||
|
|
||||||
|
if isinstance(vector, np.ndarray):
|
||||||
|
return np.array(result)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
def thick_diagonal(dim, thickness=2):
|
def thick_diagonal(dim, thickness=2):
|
||||||
row_indices = np.arange(dim).repeat(dim).reshape((dim, dim))
|
row_indices = np.arange(dim).repeat(dim).reshape((dim, dim))
|
||||||
|
Reference in New Issue
Block a user