forked from labicon/dp-ilqr
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathbbdynamicswrap.pyx
164 lines (126 loc) · 5 KB
/
bbdynamicswrap.pyx
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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
# distutils: language = c++
from enum import Enum, auto
import numpy as np
class Model(Enum):
DoubleInt4D = 0
DoubleInt6D = auto()
Car3D = auto()
Unicycle4D = auto()
Quadcopter6D = auto()
Human6D = auto()
HumanLin6D = auto()
Quadcopter12D = auto()
ctypedef void (*f_func)(double x[], double u[], double x_dot[])
cdef extern from "bbdynamics.cpp":
void rk4(
f_func dxdt, double dt, double x0[], double u[], size_t n_x, double x_new[]
)
void f_double_int_4d(double x[], double u[], double x_dot[])
void linearize_double_int_4d(double dt, double A[], double B[])
void f_double_int_6d(double x[], double u[], double x_dot[])
void linearize_double_int_6d(double dt, double A[], double B[])
void f_car_3d(double x[], double u[], double x_dot[])
void linearize_car_3d(double x[], double u[], double dt, double A[], double B[])
void f_unicycle_4d(double x[], double u[], double x_dot[])
void linearize_unicycle_4d(double x[], double u[], double dt, double A[], double B[])
void f_human_6d(double x[], double u[], double x_dot[])
void linearize_human_6d(double x[], double u[], double dt, double A[], double B[])
void f_human_lin_6d(double x[], double u[], double x_dot[])
void linearize_human_lin_6d(double dt, double A[], double B[])
void f_quad_6d(double x[], double u[], double x_dot[])
void linearize_quad_6d(double x[], double u[], double dt, double A[], double B[])
void f_quad_12d(double x[], double u[], double x_dot[])
void linearize_quad_12d(double x[], double u[], double dt, double A[], double B[])
def _common_validation(model, x, u):
if not isinstance(model, Model):
raise ValueError()
# TODO: remove when confident.
if not x.flags["C_CONTIGUOUS"] or not u.flags["C_CONTIGUOUS"]:
raise ValueError("Contiguous")
def f(x, u, model):
_common_validation(model, x, u)
cdef size_t n_x = x.shape[0]
x_dot = np.empty(n_x, dtype=np.double)
cdef double[::1] x_view = x
cdef double[::1] u_view = u
cdef double[::1] x_dot_view = x_dot
cdef f_func f
if model is Model.DoubleInt4D:
f = f_double_int_4d
elif model is Model.DoubleInt6D:
f = f_double_int_6d
elif model is Model.Car3D:
f = f_car_3d
elif model is Model.Unicycle4D:
f = f_unicycle_4d
elif model is Model.Human6D:
f = f_human_6d
elif model is Model.HumanLin6D:
f = f_human_lin_6d
elif model is Model.Quadcopter6D:
f = f_quad_6d
elif model is Model.Quadcopter12D:
f = f_quad_12d
f(&x_view[0], &u_view[0], &x_dot_view[0])
return x_dot
def integrate(x, u, double dt, model):
_common_validation(model, x, u)
cdef size_t n_x = x.shape[0]
x_new = np.empty(n_x, dtype=np.double)
cdef double[::1] x_view = x
cdef double[::1] u_view = u
cdef double[::1] x_new_view = x_new
# TODO: Figure out how to get rid of this duplicate branching logic.
cdef f_func f;
if model is Model.DoubleInt4D:
f = f_double_int_4d
elif model is Model.DoubleInt6D:
f = f_double_int_6d
elif model is Model.Car3D:
f = f_car_3d
elif model is Model.Unicycle4D:
f = f_unicycle_4d
elif model is Model.Human6D:
f = f_human_6d
elif model is Model.HumanLin6D:
f = f_human_lin_6d
elif model is Model.Quadcopter6D:
f = f_quad_6d
elif model is Model.Quadcopter12D:
f = f_quad_12d
rk4(f, dt, &x_view[0], &u_view[0], n_x, &x_new_view[0])
return x_new
def linearize(x, u, double dt, model):
if not isinstance(model, Model):
raise ValueError()
# TODO: remove when confident.
if not x.flags["C_CONTIGUOUS"]:
x = np.ascontiguousarray(x)
if not u.flags["C_CONTIGUOUS"]:
u = np.ascontiguousarray(u)
# Pre-allocate the flat output Jacobians.
cdef size_t nx = x.shape[0]
cdef size_t nu = u.shape[0]
A = np.empty((nx*nx), dtype=np.double)
B = np.empty((nx*nu), dtype=np.double)
cdef double[::1] x_view = x
cdef double[::1] u_view = u
cdef double[::1] A_view = A
cdef double[::1] B_view = B
if model is Model.DoubleInt4D:
linearize_double_int_4d(dt, &A_view[0], &B_view[0])
elif model is Model.DoubleInt6D:
linearize_double_int_6d(dt, &A_view[0], &B_view[0])
elif model is Model.Car3D:
linearize_car_3d(&x_view[0], &u_view[0], dt, &A_view[0], &B_view[0])
elif model is Model.Unicycle4D:
linearize_unicycle_4d(&x_view[0], &u_view[0], dt, &A_view[0], &B_view[0])
elif model is Model.Human6D:
linearize_human_6d(&x_view[0], &u_view[0], dt, &A_view[0], &B_view[0])
elif model is Model.HumanLin6D:
linearize_human_lin_6d(dt, &A_view[0], &B_view[0])
elif model is Model.Quadcopter6D:
linearize_quad_6d(&x_view[0], &u_view[0], dt, &A_view[0], &B_view[0])
elif model is Model.Quadcopter12D:
linearize_quad_12d(&x_view[0], &u_view[0], dt, &A_view[0], &B_view[0])
return A.reshape((nx, nx)), B.reshape((nx, nu))