-
Notifications
You must be signed in to change notification settings - Fork 1
/
math_utils.py
137 lines (102 loc) · 4.06 KB
/
math_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import math
import numpy
from pyquaternion import Quaternion
# constants
grav_const = (6.674*(10**-11)) # m^3 kg^-1 s^-2
def sign(x):
if x >= 0:
return 1
else:
return -1
# takes cartezian coords list
# gives spherical coords list
def cartesian2spherical(cart):
x = cart[0]
y = cart[1]
z = cart[2]
rho = (x**2 + y**2 + z**2)**0.5
theta = math.degrees(math.atan(((x**2 + y**2)**0.5)/z))
phi = math.degrees(math.atan(y/x))
return [rho, theta, phi]
# takes spherical coords list
# gives cartezian coords list
def spherical2cartesian(sph):
rho = sph[0]
theta = math.radians(sph[1])
phi = math.radians(sph[2])
x = math.cos(theta) * math.sin(phi) * rho
y = math.sin(theta) * math.sin(phi) * rho
z = math.cos(phi) * rho
return [x, y, z]
# cross product
def cross(a, b):
return [a[1] * b[2] - a[2] * b[1],
a[2] * b[0] - a[0] * b[2],
a[0] * b[1] - a[1] * b[0]]
# dot product
def dot(a, b):
result = 0
for a,b in zip(a,b):
result += a*b
return result
# get vector magnitude
def mag(vect):
square_sum = 0
for element in vect:
square_sum += element**2
return square_sum**0.5
# multiply vector with scalar
def vector_scale(vect, sca):
result_vec = []
for element in vect:
result_vec.append(element * sca)
return result_vec
# add vectors
def vector_add(vect1, vect2):
for i in range(len(vect1)):
vect1[i] = vect1[i] + vect2[i]
return vect1
def vector_add_safe(vect1, vect2):
result_vect = []
if len(vect1) == len(vect2):
for i in range(len(vect1)):
result_vect.append(vect1[i] + vect2[i])
else:
return -1
return result_vect
# rotate an orientation matrix
def rotate_matrix(orientation_matrix, rotation):
# orientation matrix is a 3x3 matrix, rotation is a list of three angles in degrees
orientation_matrix = numpy.array(orientation_matrix)
if rotation[0]:
rotator = Quaternion(axis=orientation_matrix[0], angle=math.radians(rotation[0]))
orientation_matrix = (numpy.array([rotator.rotate(orientation_matrix[0]), rotator.rotate(orientation_matrix[1]), rotator.rotate(orientation_matrix[2])]))
if rotation[1]:
rotator = Quaternion(axis=orientation_matrix[1], angle=math.radians(rotation[1]))
orientation_matrix = (numpy.array([rotator.rotate(orientation_matrix[0]), rotator.rotate(orientation_matrix[1]), rotator.rotate(orientation_matrix[2])]))
if rotation[2]:
rotator = Quaternion(axis=orientation_matrix[2], angle=math.radians(rotation[2]))
orientation_matrix = (numpy.array([rotator.rotate(orientation_matrix[0]), rotator.rotate(orientation_matrix[1]), rotator.rotate(orientation_matrix[2])]))
return orientation_matrix.tolist()
def abs2frame_coords(abs_coords, body):
# convert abosulte coords to body-centered reference frame coords, both cartezian
# it's like the ECEF coordinate system
return [((abs_coords[0] - body.pos[0]) * body.orient[0][0]) + ((abs_coords[1] - body.pos[1]) * body.orient[0][1]) + ((abs_coords[2] - body.pos[2]) * body.orient[0][2]),
((abs_coords[0] - body.pos[0]) * body.orient[1][0]) + ((abs_coords[1] - body.pos[1]) * body.orient[1][1]) + ((abs_coords[2] - body.pos[2]) * body.orient[1][2]),
((abs_coords[0] - body.pos[0]) * body.orient[2][0]) + ((abs_coords[1] - body.pos[1]) * body.orient[2][1]) + ((abs_coords[2] - body.pos[2]) * body.orient[2][2])]
def world2cam(w_coords, cam, factor=10):
w_coords = [-w_coords[0], -w_coords[1], -w_coords[2]]
cam_orient = cam.get_orient()
rel_pos = vector_add_safe(w_coords, vector_scale(cam.get_pos(), -1))
cam_x = cam_orient[0]
cam_y = cam_orient[1]
cam_z = cam_orient[2]
z_dist = dot(rel_pos, cam_z)
# object is behind camera, assign no position
if z_dist <= 0:
return None
x_dist = dot(rel_pos, cam_x)
y_dist = dot(rel_pos, cam_y)
x_skew = -(x_dist/z_dist) * factor
y_skew = -(y_dist/z_dist) * factor
return [x_skew, y_skew]