Trim PiecewisePolynomials (#22099)

Author: sadraddini

common/polynomial.h

@@ -175,6 +175,12 @@ class Polynomial {
 
   const std::vector<Monomial>& GetMonomials() const;
 
+  /** Returns the vector of the coefficients of the polynomial in the order
+   * of powers of the variable - the [i]th element is the coefficient of the
+   * variable raised to the ith power.
+   *
+   * @throws std::runtime_error if this Polynomial is not univariate.
+   */
   VectorX<T> GetCoefficients() const;
 
   /// Returns a set of all of the variables present in this Polynomial.

common/trajectories/BUILD.bazel

@@ -138,6 +138,7 @@ drake_cc_library(
         "//common:essential",
         "//common:name_value",
         "//common:polynomial",
+        "//math:binomial_coefficient",
         "//math:matrix_util",
         "@fmt",
     ],

common/trajectories/piecewise_polynomial.cc

@@ -11,6 +11,7 @@
 #include "drake/common/drake_assert.h"
 #include "drake/common/drake_throw.h"
 #include "drake/common/unused.h"
+#include "drake/math/binomial_coefficient.h"
 #include "drake/math/matrix_util.h"
 
 using Eigen::VectorXd;
@@ -532,6 +533,83 @@ void PiecewisePolynomial<T>::shiftRight(const T& offset) {
   }
 }
 
+namespace {
+template <typename T>
+Polynomial<T> ShiftPoly(Polynomial<T> poly, const T& x) {
+  if (poly.GetVariables().size() == 0) {
+    // Make constant polynomial.
+    return poly;
+  }
+  using std::pow;
+  // Given p(t) = a0 + a1*t + a2*t^2 + ... + an*t^n,
+  // we want to shift the parameter to p(t + x)
+  // = b0 + b1*t + b2*t^2 + ... + bn*t^n.
+  // We can expand the rhs to get the values in b.
+  DRAKE_THROW_UNLESS(poly.GetVariables().size() == 1);
+  int n = poly.GetDegree();
+  const auto& a = poly.GetCoefficients();
+  Eigen::Matrix<T, Eigen::Dynamic, 1> b =
+      Eigen::Matrix<T, Eigen::Dynamic, 1>::Zero(n + 1);
+  for (int i = 0; i <= n; i++) {
+    for (int j = 0; j <= i; j++) {
+      b(j) += a(i) * drake::math::BinomialCoefficient(i, j) * pow(x, i - j);
+    }
+  }
+  return Polynomial<T>(b);
+}
+}  // namespace
+
+template <typename T>
+int PiecewisePolynomial<T>::AddBreak(const T& new_break) {
+  auto& breaks = this->get_mutable_breaks();
+  // Check if the new break is already within the kEpsilonTime of an
+  // existing break.
+  for (int k = 0; k < this->get_number_of_segments(); k++) {
+    const T& break_time = breaks[k];
+    if (abs(break_time - new_break) < PiecewiseTrajectory<T>::kEpsilonTime) {
+      // No need to add a new break.
+      return k;
+    }
+  }
+  DRAKE_THROW_UNLESS(this->is_time_in_range(new_break));
+  const auto value = this->value(new_break);
+  int i = this->get_segment_index(new_break);
+  const T last_break_before_new = this->start_time(i);
+  const T shift = new_break - last_break_before_new;
+  breaks.insert(breaks.begin() + i + 1, new_break);
+  // New polynomial matrix to be added at index i.
+  const auto& broken_polynomial = polynomials_[i];
+  PolynomialMatrix new_polynomial = broken_polynomial;
+  for (int row = 0; row < rows(); row++) {
+    for (int col = 0; col < cols(); col++) {
+      new_polynomial(row, col) = ShiftPoly(new_polynomial(row, col), shift);
+    }
+  }
+  // The i+1'th polynomial should become the new_polynomial.
+  polynomials_.insert(polynomials_.begin() + i + 1, new_polynomial);
+  return i + 1;
+}
+
+template <typename T>
+PiecewisePolynomial<T> PiecewisePolynomial<T>::SliceByTime(
+    const T& start_time, const T& end_time) const {
+  DRAKE_THROW_UNLESS(start_time < end_time);
+  DRAKE_THROW_UNLESS(this->is_time_in_range(start_time));
+  DRAKE_THROW_UNLESS(this->is_time_in_range(end_time));
+  PiecewisePolynomial<T> result{*this};
+  int i = result.AddBreak(start_time);
+  int j = result.AddBreak(end_time);
+  std::vector<PolynomialMatrix> polynomials;
+  std::vector<T> breaks;
+  for (int k = i; k < j; k++) {
+    polynomials.push_back(result.getPolynomialMatrix(k));
+    breaks.push_back(result.get_segment_times()[k]);
+  }
+  // Add the last break
+  breaks.push_back(result.get_segment_times()[j]);
+  return PiecewisePolynomial<T>(polynomials, breaks);
+}
+
 template <typename T>
 void PiecewisePolynomial<T>::setPolynomialMatrixBlock(
     const typename PiecewisePolynomial<T>::PolynomialMatrix& replacement,

common/trajectories/piecewise_polynomial.h

@@ -745,6 +745,26 @@ class PiecewisePolynomial final : public PiecewiseTrajectory<T> {
    */
   void shiftRight(const T& offset);
 
+  /** Adds a break at the specified time. It does not change the value of the
+   * trajectory at any point but the number of segments increases by 1.
+   * @returns the index of the new break.
+   * @throws std::exception if `new_break` is not within the trajectory's
+   *        time range.
+   * @warning If `new_break` is within PiecewiseTrajectory::kEpsilonTime from
+   *         an existing break, the new break will be silently ignored. Returns
+   *        the index of the existing break.
+   *
+   */
+  int AddBreak(const T& new_break);
+
+  /** Slices the trajectory within a specified time range.
+   * q = p.SliceByTime(t1, t2) returns a PiecewisePolynomial q such that
+   * q.start_time() = t1, q.end_time() = t2, and q(t) = p(t) for t1 <= t <= t2.
+   * @throws std::exception if `start_time` or `end_time` is not within the
+   *        trajectory's time range.
+   */
+  PiecewisePolynomial SliceByTime(const T& start_time, const T& end_time) const;
+
   /**
    * Replaces the specified block of the PolynomialMatrix at the given
    * segment index.

common/trajectories/test/piecewise_polynomial_test.cc

@@ -147,6 +147,30 @@ void testBasicFunctionality() {
                                               piecewise1.end_time());
     double t = uniform(generator);
 
+    // Random segment index to put the break
+    uniform_int_distribution<> segment_dist(
+        0, piecewise2.get_number_of_segments() - 1);
+    int segment_index = segment_dist(generator);
+
+    PiecewisePolynomial<T> piecewise1_with_new_break = piecewise1;
+    EXPECT_EQ(piecewise1_with_new_break.AddBreak(t),
+              piecewise1.get_segment_index(t) + 1);
+    EXPECT_EQ(piecewise1_with_new_break.get_number_of_segments(),
+              piecewise1.get_number_of_segments() + 1);
+
+    PiecewisePolynomial<T> piecewise2_with_new_break = piecewise2;
+    EXPECT_EQ(piecewise2_with_new_break.AddBreak(
+                  piecewise2.start_time(segment_index)),
+              segment_index);
+
+    EXPECT_EQ(piecewise2_with_new_break.get_number_of_segments(),
+              piecewise2.get_number_of_segments());
+
+    double t_2 = uniform(generator);
+
+    auto trimmed_piecewise2 =
+        piecewise2.SliceByTime(std::min(t, t_2), std::max(t, t_2));
+
     EXPECT_TRUE(CompareMatrices(sum.value(t),
                                 piecewise1.value(t) + piecewise2.value(t), 1e-8,
                                 MatrixCompareType::absolute));
@@ -183,6 +207,31 @@ void testBasicFunctionality() {
     EXPECT_TRUE(CompareMatrices(piecewise2_twice.value(t),
                                 piecewise2_twice.value(t + total_time), 1e-8,
                                 MatrixCompareType::absolute));
+
+    // The piecewise_with_new_break values must not change.
+    EXPECT_TRUE(CompareMatrices(piecewise1_with_new_break.value(t),
+                                piecewise1.value(t), 1e-8,
+                                MatrixCompareType::absolute));
+    // Check the samples at the middle of every segment.
+    for (int k = 0; k < piecewise2.get_number_of_segments(); ++k) {
+      double t_middle = (piecewise2.start_time(k) + piecewise2.end_time(k)) / 2;
+      EXPECT_TRUE(CompareMatrices(piecewise2.value(t_middle),
+                                  piecewise2_with_new_break.value(t_middle),
+                                  1e-8, MatrixCompareType::absolute));
+    }
+
+    // Check if trimmed_piecewise2 is indeed a trimmed version of piecewise2.
+    EXPECT_EQ(trimmed_piecewise2.start_time(), std::min(t, t_2));
+    EXPECT_EQ(trimmed_piecewise2.end_time(), std::max(t, t_2));
+    uniform_real_distribution<double> uniform_trimmed(
+        trimmed_piecewise2.start_time(), trimmed_piecewise2.end_time());
+    // Pick 10 random samples and check that the sampled values are the same.
+    for (int k = 0; k < 10; k++) {
+      double t_sample = uniform_trimmed(generator);
+      EXPECT_TRUE(CompareMatrices(trimmed_piecewise2.value(t_sample),
+                                  piecewise2.value(t_sample), 1e-8,
+                                  MatrixCompareType::absolute));
+    }
   }
 }