class PPTrajectory():
def __init__(self, sample_times, num_vars, degree, continuity_degree):
self.sample_times = sample_times
self.n = num_vars
self.degree = degree
self.prog = MathematicalProgram()
self.coeffs = []
for i in range(len(sample_times)):
self.coeffs.append(
self.prog.NewContinuousVariables(num_vars, degree + 1, "C"))
self.result = None
# Add continuity constraints
for s in range(len(sample_times) - 1):
trel = sample_times[s + 1] - sample_times[s]
coeffs = self.coeffs[s]
for var in range(self.n):
for deg in range(continuity_degree + 1):
# Don't use eval here, because I want left and right
# values of the same time
left_val = 0
for d in range(deg, self.degree + 1):
left_val += coeffs[var, d]*np.power(trel, d-deg) * \
math.factorial(d)/math.factorial(d-deg)
right_val = self.coeffs[s + 1][var,
deg] * math.factorial(deg)
self.prog.AddLinearConstraint(left_val == right_val)
# Add cost to minimize highest order terms
for s in range(len(sample_times) - 1):
self.prog.AddQuadraticCost(np.eye(num_vars), np.zeros(
(num_vars, 1)), self.coeffs[s][:, -1])
def eval(self, t, derivative_order=0):
if derivative_order > self.degree:
return 0
s = 0
while s < len(self.sample_times) - 1 and t >= self.sample_times[s + 1]:
s += 1
trel = t - self.sample_times[s]
if self.result is None:
coeffs = self.coeffs[s]
else:
coeffs = self.result.GetSolution(self.coeffs[s])
deg = derivative_order
val = 0 * coeffs[:, 0]
for var in range(self.n):
for d in range(deg, self.degree + 1):
val[var] += coeffs[var, d]*np.power(trel, d-deg) * \
math.factorial(d)/math.factorial(d-deg)
return val
def add_constraint(self, t, derivative_order, lb, ub=None):
'''
Adds a constraint of the form d^deg lb <= x(t) / dt^deg <= ub
'''
if ub is None:
ub = lb
assert (derivative_order <= self.degree)
val = self.eval(t, derivative_order)
self.prog.AddLinearConstraint(val, lb, ub)
def generate(self):
self.result = Solve(self.prog)
assert (self.result.is_success())
def plot_sublevelset_expression(ax, e, vertices=51, **kwargs):
"""
Plots the 2D sub-level set e(x) <= 1, which must contain the origin.
Args:
ax: the matplotlib axis to receive the plot
e: a symbolic expression in two variables
vertices: number of sample points along the boundary
kwargs: are passed to the matplotlib fill method
Returns:
the return values from matplotlib's fill command.
"""
x = list(e.GetVariables())
assert len(x) == 2, "e must be an expression in two variables"
# Handle the special case where e is a degree 2 polynomial.
if e.is_polynomial():
p = Polynomial(e)
if p.TotalDegree() == 2:
env = {a: 0 for a in x}
c = e.Evaluate(env)
e1 = e.Jacobian(x)
b = Evaluate(e1, env)
e2 = Jacobian(e1, x)
A = 0.5 * Evaluate(e2, env)
return plot_sublevelset_quadratic(ax, A, b, c, vertices, **kwargs)
# Find the level-set in polar coordinates, by sampling theta and
# root-finding (on the scalar expression) to find a rplus and rminus.
Xplus = np.empty((2, vertices))
Xminus = np.empty((2, vertices))
i = 0
for theta in np.linspace(0, np.pi, vertices):
prog = MathematicalProgram()
r = prog.NewContinuousVariables(1, "r")[0]
env = {x[0]: r * np.cos(theta), x[1]: r * np.sin(theta)}
scalar = e.Substitute(env)
b = prog.AddBoundingBoxConstraint(0, np.inf, r)
prog.AddConstraint(scalar == 1)
prog.AddQuadraticCost([1], [0], [r])
prog.SetInitialGuess(r, 0.1) # or anything non-zero.
result = Solve(prog)
assert result.is_success(), "Failed to find the level set"
rplus = result.GetSolution(r)
Xplus[0, i] = rplus * np.cos(theta)
Xplus[1, i] = rplus * np.sin(theta)
b.evaluator().UpdateLowerBound([-np.inf])
b.evaluator().UpdateUpperBound([0])
prog.SetInitialGuess(r, -0.1) # or anything non-zero.
result = Solve(prog)
assert result.is_success(), "Failed to find the level set"
rminus = result.GetSolution(r)
Xminus[0, i] = rminus * np.cos(theta)
Xminus[1, i] = rminus * np.sin(theta)
i = i + 1
return ax.fill(np.hstack((Xplus[0, :], Xminus[0, :])),
np.hstack((Xplus[1, :], Xminus[1, :])), **kwargs)