Skip to content

Commit

Permalink
Switch to sparse solver for CubicWithContinuousSecondDerivatives
Browse files Browse the repository at this point in the history
  • Loading branch information
@mpetersen94
mpetersen94 committed Jun 28, 2021
1 parent dad3e0c commit 4550c8c
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 50 deletions.
134 changes: 85 additions & 49 deletions common/trajectories/piecewise_polynomial.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
#include <memory>
#include <utility>

#include <Eigen/SparseCore>
#include <Eigen/SparseLU>
#include <fmt/format.h>

#include "drake/common/drake_assert.h"
Expand Down Expand Up @@ -810,20 +812,18 @@ int PiecewisePolynomial<T>::
const std::vector<T>& breaks,
const std::vector<MatrixX<T>>& samples,
int row, int col,
MatrixX<T>* A,
std::vector<Eigen::Triplet<T>>* triplet_list,
VectorX<T>* b) {
const std::vector<T>& times = breaks;
const std::vector<MatrixX<T>>& Y = samples;
int N = static_cast<int>(times.size());

DRAKE_DEMAND(A != nullptr);
DRAKE_DEMAND(triplet_list != nullptr);
DRAKE_DEMAND(b != nullptr);
DRAKE_DEMAND(A->rows() == 3 * (N - 1));
DRAKE_DEMAND(A->cols() == 3 * (N - 1));
DRAKE_DEMAND(b->rows() == 3 * (N - 1));

int row_idx = 0;
MatrixX<T>& Aref = *A;
std::vector<Eigen::Triplet<T>>& triplet_ref = *triplet_list;
VectorX<T>& bref = *b;

for (int i = 0; i < N - 1; ++i) {
Expand All @@ -832,26 +832,26 @@ int PiecewisePolynomial<T>::
// y_i(x_{i+1}) = y_{i+1}(x_{i}) =>
// Y[i] + a1i*(x_{i+1} - x_i) + a2i(x_{i+1} - x_i)^2 + a3i(x_{i+1} -
// x_i)^3 = Y[i+1]
Aref(row_idx, 3 * i + 0) = dt;
Aref(row_idx, 3 * i + 1) = dt * dt;
Aref(row_idx, 3 * i + 2) = dt * dt * dt;
bref(row_idx++) = Y[i + 1](row, col)-Y[i](row, col);
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 0, dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 1, dt * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 2, dt * dt * dt));
bref(row_idx++) = Y[i + 1](row, col) - Y[i](row, col);

// y_i'(x_{i+1}) = y_{i+1}'(x_{i}) =>
// a1i + 2*a2i(x_{i+1} - x_i) + 3*a3i(x_{i+1} - x_i)^2 = a1{i+1}
if (i != N - 2) {
Aref(row_idx, 3 * i + 0) = 1;
Aref(row_idx, 3 * i + 1) = 2 * dt;
Aref(row_idx, 3 * i + 2) = 3 * dt * dt;
Aref(row_idx++, 3 * (i + 1)) = -1;
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 0, 1));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 1, 2 * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 2, 3 * dt * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx++, 3 * (i + 1), -1));
}

if (i != N - 2) {
// y_i''(x_{i+1}) = y_{i+1}''(x_{i}) =>
// 2*a2i + 6*a3i(x_{i+1} - x_i) = 2*a2{i+1}
Aref(row_idx, 3 * i + 1) = 2;
Aref(row_idx, 3 * i + 2) = 6 * dt;
Aref(row_idx++, 3 * (i + 1) + 1) = -2;
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 1, 2));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx, 3 * i + 2, 6 * dt));
triplet_ref.push_back(Eigen::Triplet<T>(row_idx++, 3 * (i + 1) + 1, -2));
}
}
DRAKE_DEMAND(row_idx == 3 * (N - 1) - 2);
Expand Down Expand Up @@ -893,32 +893,45 @@ PiecewisePolynomial<T>::CubicWithContinuousSecondDerivatives(
polynomials[i].resize(rows, cols);
}

MatrixX<T> A(3 * (N - 1), 3 * (N - 1));
Eigen::SparseMatrix<T> A(3 * (N - 1), 3 * (N - 1));
VectorX<T> b(3 * (N - 1));
VectorX<T> solution;
VectorX<T> coeffs(4);
Eigen::SparseLU<Eigen::SparseMatrix<T>> solver;

A.setZero();
b.setZero();

// Sets up a linear equation to solve for the coefficients.
for (int j = 0; j < rows; ++j) {
for (int k = 0; k < cols; ++k) {
int row_idx =
SetupCubicSplineInteriorCoeffsLinearSystem(times, Y, j, k, &A, &b);
std::vector<Eigen::Triplet<T>> triplet_list;
// 10 coefficients are needed for the constraints that ensure continuity
// between adjacent segments: 3 for position, 4 for velocity, 3 for
// acceleration. The additional coefficients consist of: 3 for final
// position constraint, 1 for initial velocity constraint & 3 for final
// velocity constraint.
triplet_list.reserve(10 * (N - 2) + 7);
int row_idx = SetupCubicSplineInteriorCoeffsLinearSystem(
times, Y, j, k, &triplet_list, &b);

// Endpoints' velocity matches the given ones.
A(row_idx, 0) = 1;
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 0, 1));
b(row_idx++) = Ydot_start(j, k);

A(row_idx, 3 * (N - 2) + 0) = 1;
A(row_idx, 3 * (N - 2) + 1) = 2 * (times[N - 1] - times[N - 2]);
A(row_idx, 3 * (N - 2) + 2) =
3 * (times[N - 1] - times[N - 2]) * (times[N - 1] - times[N - 2]);
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 0, 1));
triplet_list.push_back(Eigen::Triplet<T>(
row_idx, 3 * (N - 2) + 1, 2 * (times[N - 1] - times[N - 2])));
triplet_list.push_back(Eigen::Triplet<T>(
row_idx, 3 * (N - 2) + 2,
3 * (times[N - 1] - times[N - 2]) * (times[N - 1] - times[N - 2])));
b(row_idx++) = Ydot_end(j, k);

// TODO(siyuan.feng): Should switch to a sparse solver.
solution = A.partialPivLu().solve(b);
A.setFromTriplets(triplet_list.begin(), triplet_list.end());
if (j == 0 && k == 0) {
solver.analyzePattern(A);
}
solver.factorize(A);
solution = solver.solve(b);

for (int i = 0; i < N - 1; ++i) {
coeffs(0) = Y[i](j, k);
Expand Down Expand Up @@ -957,63 +970,86 @@ PiecewisePolynomial<T>::CubicWithContinuousSecondDerivatives(
polynomials[i].resize(rows, cols);
}

MatrixX<T> A(3 * (N - 1), 3 * (N - 1));
Eigen::SparseMatrix<T> A(3 * (N - 1), 3 * (N - 1));
VectorX<T> b(3 * (N - 1));
VectorX<T> solution;
VectorX<T> coeffs(4);
Eigen::SparseLU<Eigen::SparseMatrix<T>> solver;

A.setZero();
b.setZero();

// Sets up a linear equation to solve for the coefficients.
for (int j = 0; j < rows; ++j) {
for (int k = 0; k < cols; ++k) {
int row_idx =
SetupCubicSplineInteriorCoeffsLinearSystem(times, Y, j, k, &A, &b);
std::vector<Eigen::Triplet<T>> triplet_list;
// 10 coefficients are needed for the constraints that ensure continuity
// between adjacent segments: 3 for position, 4 for velocity, 3 for
// acceleration. The additional coefficients consist of: 3 for final
// position constraint, 4 for periodic velocity constraint & 3 for
// periodic acceleration constraint. In the case of "not-a-sample" end
// condition, fewer coefficients are needed.
triplet_list.reserve(10 * (N - 2) + 10);
int row_idx = SetupCubicSplineInteriorCoeffsLinearSystem(
times, Y, j, k, &triplet_list, &b);

if (periodic_end_condition) {
// Time during the last segment.
const T end_dt = times[times.size() - 1] - times[times.size() - 2];
// Enforce velocity between end-of-last and beginning-of-first segments
// is continuous.
A(row_idx, 0) = -1; // Linear term of 1st segment.
A(row_idx, 3 * (N - 2) + 0) = 1; // Linear term of last segment.
A(row_idx, 3 * (N - 2) + 1) = 2 * end_dt; // Squared term of last
// segment.
A(row_idx, 3 * (N - 2) + 2) = 3 * end_dt * end_dt; // Cubic term of
// last segment.
// Linear term of 1st segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 0, -1));
// Linear term of last segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 0, 1));
// Squared term of last segment.
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 1, 2 * end_dt));
// Cubic term of last segment.
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 2, 3 * end_dt * end_dt));
b(row_idx++) = 0;

// Enforce that acceleration between end-of-last and beginning-of-first
// segments is continuous.
A(row_idx, 1) = -2; // Quadratic term of 1st segment.
A(row_idx, 3 * (N - 2) + 1) = 2; // Quadratic term of last segment.
A(row_idx, 3 * (N - 2) + 2) = 6 * end_dt; // Cubic term of last
// segment.
// Quadratic term of 1st segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 1, -2));
// Quadratic term of last segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 1, 2));
// Cubic term of last segment.
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 2, 6 * end_dt));
b(row_idx++) = 0;
} else {
if (N > 3) {
// Ydddot(times[1]) is continuous.
A(row_idx, 2) = 1; // Cubic term of 1st segment.
A(row_idx, 3 + 2) = -1; // Cubic term of 2nd segment.
// Cubic term of 1st segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 2, 1));
// Cubic term of 2nd segment.
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 + 2, -1));
b(row_idx++) = 0;

// Ydddot(times[N-2]) is continuous.
A(row_idx, 3 * (N - 3) + 2) = 1;
A(row_idx, 3 * (N - 2) + 2) = -1;
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 3) + 2, 1));
triplet_list.push_back(
Eigen::Triplet<T>(row_idx, 3 * (N - 2) + 2, -1));
b(row_idx++) = 0;
} else {
// Set Jerk to zero if only have 3 points, becomes a quadratic.
A(row_idx, 2) = 1;
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 2, 1));
b(row_idx++) = 0;

A(row_idx, 3 + 2) = 1;
triplet_list.push_back(Eigen::Triplet<T>(row_idx, 3 + 2, 1));
b(row_idx++) = 0;
}
}

// TODO(siyuan.feng): Should switch to a sparse solver.
solution = A.partialPivLu().solve(b);
A.setFromTriplets(triplet_list.begin(), triplet_list.end());
if (j == 0 && k == 0) {
solver.analyzePattern(A);
}
solver.factorize(A);
solution = solver.solve(b);

for (int i = 0; i < N - 1; ++i) {
coeffs(0) = Y[i](j, k);
Expand Down
2 changes: 1 addition & 1 deletion common/trajectories/piecewise_polynomial.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -822,7 +822,7 @@ class PiecewisePolynomial final : public PiecewiseTrajectory<T> {
static int SetupCubicSplineInteriorCoeffsLinearSystem(
const std::vector<T>& breaks,
const std::vector<MatrixX<T>>& samples, int row, int col,
MatrixX<T>* A, VectorX<T>* b);
std::vector<Eigen::Triplet<T>>* triplet_list, VectorX<T>* b);

// Throws std::runtime_error if
// `breaks` and `samples` have different length,
Expand Down

0 comments on commit 4550c8c

Please sign in to comment.