Add transform methods to Trajectory (#2187)

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,