Have rotate_vector return same type as input

This commit is contained in:
Grant Sanderson
2020-03-06 06:15:28 -08:00
parent 8fc6c8107a
commit cfe72922ca

View File

@ -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))