Add transform methods to Trajectory (#2187)

wpilibc/src/main/native/cpp/trajectory/Trajectory.cpp

@@ -107,6 +107,33 @@ Trajectory::State Trajectory::Sample(units::second_t t) const {
                                 (t - prevSample.t) / (sample.t - prevSample.t));
 }
 
+Trajectory Trajectory::TransformBy(const Transform2d& transform) {
+  auto& firstState = m_states[0];
+  auto& firstPose = firstState.pose;
+
+  // Calculate the transformed first pose.
+  auto newFirstPose = firstPose + transform;
+  auto newStates = m_states;
+  newStates[0].pose = newFirstPose;
+
+  for (unsigned int i = 1; i < newStates.size(); i++) {
+    auto& state = newStates[i];
+    // We are transforming relative to the coordinate frame of the new initial
+    // pose.
+    state.pose = newFirstPose + (state.pose - firstPose);
+  }
+
+  return Trajectory(newStates);
+}
+
+Trajectory Trajectory::RelativeTo(const Pose2d& pose) {
+  auto newStates = m_states;
+  for (auto& state : newStates) {
+    state.pose = state.pose.RelativeTo(pose);
+  }
+  return Trajectory(newStates);
+}
+
 void frc::to_json(wpi::json& json, const Trajectory::State& state) {
   json = wpi::json{{"time", state.t.to<double>()},
                    {"velocity", state.velocity.to<double>()},

wpilibc/src/main/native/include/frc/trajectory/Trajectory.h

@@ -12,6 +12,7 @@
 #include <units/units.h>
 
 #include "frc/geometry/Pose2d.h"
+#include "frc/geometry/Transform2d.h"
 
 namespace wpi {
 class json;
@@ -105,6 +106,27 @@ class Trajectory {
    */
   State Sample(units::second_t t) const;
 
+  /**
+   * Transforms all poses in the trajectory by the given transform. This is
+   * useful for converting a robot-relative trajectory into a field-relative
+   * trajectory. This works with respect to the first pose in the trajectory.
+   *
+   * @param transform The transform to transform the trajectory by.
+   * @return The transformed trajectory.
+   */
+  Trajectory TransformBy(const Transform2d& transform);
+
+  /**
+   * Transforms all poses in the trajectory so that they are relative to the
+   * given pose. This is useful for converting a field-relative trajectory
+   * into a robot-relative trajectory.
+   *
+   * @param pose The pose that is the origin of the coordinate frame that
+   *             the current trajectory will be transformed into.
+   * @return The transformed trajectory.
+   */
+  Trajectory RelativeTo(const Pose2d& pose);
+
   /**
    * Returns the initial pose of the trajectory.
    *

wpilibc/src/test/native/cpp/trajectory/TrajectoryTransformTest.cpp

@@ -0,0 +1,70 @@
+/*----------------------------------------------------------------------------*/
+/* Copyright (c) 2019 FIRST. All Rights Reserved.                             */
+/* Open Source Software - may be modified and shared by FRC teams. The code   */
+/* must be accompanied by the FIRST BSD license file in the root directory of */
+/* the project.                                                               */
+/*----------------------------------------------------------------------------*/
+
+#include <vector>
+
+#include "frc/trajectory/Trajectory.h"
+#include "frc/trajectory/TrajectoryConfig.h"
+#include "frc/trajectory/TrajectoryGenerator.h"
+#include "gtest/gtest.h"
+
+void TestSameShapedTrajectory(std::vector<frc::Trajectory::State> statesA,
+                              std::vector<frc::Trajectory::State> statesB) {
+  for (unsigned int i = 0; i < statesA.size() - 1; i++) {
+    auto a1 = statesA[i].pose;
+    auto a2 = statesA[i + 1].pose;
+
+    auto b1 = statesB[i].pose;
+    auto b2 = statesB[i + 1].pose;
+
+    auto a = a2.RelativeTo(a1);
+    auto b = b2.RelativeTo(b1);
+
+    EXPECT_NEAR(a.Translation().X().to<double>(),
+                b.Translation().X().to<double>(), 1E-9);
+    EXPECT_NEAR(a.Translation().Y().to<double>(),
+                b.Translation().Y().to<double>(), 1E-9);
+    EXPECT_NEAR(a.Rotation().Radians().to<double>(),
+                b.Rotation().Radians().to<double>(), 1E-9);
+  }
+}
+
+TEST(TrajectoryTransforms, TransformBy) {
+  frc::TrajectoryConfig config{3_mps, 3_mps_sq};
+  auto trajectory = frc::TrajectoryGenerator::GenerateTrajectory(
+      frc::Pose2d{}, {}, frc::Pose2d{1_m, 1_m, frc::Rotation2d(90_deg)},
+      config);
+
+  auto transformedTrajectory =
+      trajectory.TransformBy({{1_m, 2_m}, frc::Rotation2d(30_deg)});
+
+  auto firstPose = transformedTrajectory.Sample(0_s).pose;
+
+  EXPECT_NEAR(firstPose.Translation().X().to<double>(), 1.0, 1E-9);
+  EXPECT_NEAR(firstPose.Translation().Y().to<double>(), 2.0, 1E-9);
+  EXPECT_NEAR(firstPose.Rotation().Degrees().to<double>(), 30.0, 1E-9);
+
+  TestSameShapedTrajectory(trajectory.States(), transformedTrajectory.States());
+}
+
+TEST(TrajectoryTransforms, RelativeTo) {
+  frc::TrajectoryConfig config{3_mps, 3_mps_sq};
+  auto trajectory = frc::TrajectoryGenerator::GenerateTrajectory(
+      frc::Pose2d{1_m, 2_m, frc::Rotation2d(30_deg)}, {},
+      frc::Pose2d{5_m, 7_m, frc::Rotation2d(90_deg)}, config);
+
+  auto transformedTrajectory =
+      trajectory.RelativeTo({1_m, 2_m, frc::Rotation2d(30_deg)});
+
+  auto firstPose = transformedTrajectory.Sample(0_s).pose;
+
+  EXPECT_NEAR(firstPose.Translation().X().to<double>(), 0, 1E-9);
+  EXPECT_NEAR(firstPose.Translation().Y().to<double>(), 0, 1E-9);
+  EXPECT_NEAR(firstPose.Rotation().Degrees().to<double>(), 0, 1E-9);
+
+  TestSameShapedTrajectory(trajectory.States(), transformedTrajectory.States());
+}

wpilibj/src/main/java/edu/wpi/first/wpilibj/trajectory/Trajectory.java

@@ -7,6 +7,7 @@
 
 package edu.wpi.first.wpilibj.trajectory;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 import java.util.stream.Collectors;
@@ -14,6 +15,7 @@ import java.util.stream.Collectors;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import edu.wpi.first.wpilibj.geometry.Pose2d;
+import edu.wpi.first.wpilibj.geometry.Transform2d;
 
 /**
  * Represents a time-parameterized trajectory. The trajectory contains of
@@ -141,6 +143,57 @@ public class Trajectory {
         (timeSeconds - prevSample.timeSeconds) / (sample.timeSeconds - prevSample.timeSeconds));
   }
 
+  /**
+   * Transforms all poses in the trajectory by the given transform. This is
+   * useful for converting a robot-relative trajectory into a field-relative
+   * trajectory. This works with respect to the first pose in the trajectory.
+   *
+   * @param transform The transform to transform the trajectory by.
+   * @return The transformed trajectory.
+   */
+  @SuppressWarnings("PMD.AvoidInstantiatingObjectsInLoops")
+  public Trajectory transformBy(Transform2d transform) {
+    var firstState = m_states.get(0);
+    var firstPose = firstState.poseMeters;
+
+    // Calculate the transformed first pose.
+    var newFirstPose = firstPose.plus(transform);
+    List<State> newStates = new ArrayList<>();
+
+    newStates.add(new State(
+        firstState.timeSeconds, firstState.velocityMetersPerSecond,
+        firstState.accelerationMetersPerSecondSq, newFirstPose, firstState.curvatureRadPerMeter
+    ));
+
+    for (int i = 1; i < m_states.size(); i++) {
+      var state = m_states.get(i);
+      // We are transforming relative to the coordinate frame of the new initial pose.
+      newStates.add(new State(
+          state.timeSeconds, state.velocityMetersPerSecond,
+          state.accelerationMetersPerSecondSq, newFirstPose.plus(state.poseMeters.minus(firstPose)),
+          state.curvatureRadPerMeter
+      ));
+    }
+
+    return new Trajectory(newStates);
+  }
+
+  /**
+   * Transforms all poses in the trajectory so that they are relative to the
+   * given pose. This is useful for converting a field-relative trajectory
+   * into a robot-relative trajectory.
+   *
+   * @param pose The pose that is the origin of the coordinate frame that
+   *             the current trajectory will be transformed into.
+   * @return The transformed trajectory.
+   */
+  public Trajectory relativeTo(Pose2d pose) {
+    return new Trajectory(m_states.stream().map(state -> new State(state.timeSeconds,
+        state.velocityMetersPerSecond, state.accelerationMetersPerSecondSq,
+        state.poseMeters.relativeTo(pose), state.curvatureRadPerMeter))
+        .collect(Collectors.toList()));
+  }
+
   /**
    * Represents a time-parameterized trajectory. The trajectory contains of
    * various States that represent the pose, curvature, time elapsed, velocity,

wpilibj/src/test/java/edu/wpi/first/wpilibj/trajectory/TrajectoryTransformTest.java

@@ -0,0 +1,68 @@
+/*----------------------------------------------------------------------------*/
+/* Copyright (c) 2019 FIRST. All Rights Reserved.                             */
+/* Open Source Software - may be modified and shared by FRC teams. The code   */
+/* must be accompanied by the FIRST BSD license file in the root directory of */
+/* the project.                                                               */
+/*----------------------------------------------------------------------------*/
+
+package edu.wpi.first.wpilibj.trajectory;
+
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+
+import edu.wpi.first.wpilibj.geometry.Pose2d;
+import edu.wpi.first.wpilibj.geometry.Rotation2d;
+import edu.wpi.first.wpilibj.geometry.Transform2d;
+import edu.wpi.first.wpilibj.geometry.Translation2d;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+class TrajectoryTransformTest {
+  @Test
+  void testTransformBy() {
+    var config = new TrajectoryConfig(3, 3);
+    var trajectory = TrajectoryGenerator.generateTrajectory(
+        new Pose2d(), List.of(), new Pose2d(1, 1, Rotation2d.fromDegrees(90)),
+        config
+    );
+
+    var transformedTrajectory = trajectory.transformBy(
+        new Transform2d(new Translation2d(1, 2), Rotation2d.fromDegrees(30)));
+
+    // Test initial pose.
+    assertEquals(new Pose2d(1, 2, Rotation2d.fromDegrees(30)),
+        transformedTrajectory.sample(0).poseMeters);
+
+    testSameShapedTrajectory(trajectory.getStates(), transformedTrajectory.getStates());
+  }
+
+  @Test
+  void testRelativeTo() {
+    var config = new TrajectoryConfig(3, 3);
+    var trajectory = TrajectoryGenerator.generateTrajectory(
+        new Pose2d(1, 2, Rotation2d.fromDegrees(30.0)),
+        List.of(), new Pose2d(5, 7, Rotation2d.fromDegrees(90)),
+        config
+    );
+
+    var transformedTrajectory = trajectory.relativeTo(new Pose2d(1, 2, Rotation2d.fromDegrees(30)));
+
+    // Test initial pose.
+    assertEquals(new Pose2d(), transformedTrajectory.sample(0).poseMeters);
+
+    testSameShapedTrajectory(trajectory.getStates(), transformedTrajectory.getStates());
+  }
+
+  void testSameShapedTrajectory(List<Trajectory.State> statesA, List<Trajectory.State> statesB) {
+    for (int i = 0; i < statesA.size() - 1; i++) {
+      var a1 = statesA.get(i).poseMeters;
+      var a2 = statesA.get(i + 1).poseMeters;
+
+      var b1 = statesB.get(i).poseMeters;
+      var b2 = statesB.get(i + 1).poseMeters;
+
+      assertEquals(a2.relativeTo(a1), b2.relativeTo(b1));
+    }
+  }
+}