Merge OptimalTrajectoryGenerator and SwerveDiscreteOptimal classes

The geometry classes now overload operator== to produce equality
constraints.

build.rs

@@ -23,6 +23,7 @@ fn main() {
         .file("src/trajoptlibrust.cpp")
         .include("src")
         .include(format!("{}/include", cmake_dest.display()))
+        .include(format!("{}/include/eigen3", cmake_dest.display()))
         .std("c++20");
 
     bridge_build.compile("trajoptrust");

examples/Swerve/src/Main.cpp

@@ -1,6 +1,6 @@
 // Copyright (c) TrajoptLib contributors
 
-#include <trajopt/OptimalTrajectoryGenerator.hpp>
+#include <trajopt/SwerveTrajectoryGenerator.hpp>
 
 int main() {
   trajopt::SwerveDrivetrain swerveDrivetrain{
@@ -23,7 +23,9 @@ int main() {
   path.WptZeroVelocity(1);
   path.ControlIntervalCounts({4});
 
+  trajopt::SwerveTrajectoryGenerator generator{path};
+
   // SOLVE
   [[maybe_unused]]
-  auto solution = trajopt::OptimalTrajectoryGenerator::Generate(path, true);
+  auto solution = generator.Generate(true);
 }

include/trajopt/SwerveTrajectoryGenerator.hpp

@@ -0,0 +1,82 @@
+// Copyright (c) TrajoptLib contributors
+
+#pragma once
+
+#include <stdint.h>
+
+#include <functional>
+#include <string>
+#include <vector>
+
+#include <sleipnir/optimization/OptimizationProblem.hpp>
+
+#include "trajopt/path/SwervePathBuilder.hpp"
+#include "trajopt/solution/SwerveSolution.hpp"
+#include "trajopt/util/SymbolExports.hpp"
+#include "trajopt/util/expected"
+
+namespace trajopt {
+
+/**
+ * This trajectory generator class contains functions to generate
+ * time-optimal trajectories for several drivetrain types.
+ */
+class TRAJOPT_DLLEXPORT SwerveTrajectoryGenerator {
+ public:
+  /**
+   * Construct a new swerve trajectory optimization problem.
+   *
+   * @param pathBuilder The path builder.
+   * @param handle An identifier for state callbacks.
+   */
+  explicit SwerveTrajectoryGenerator(const SwervePathBuilder& pathBuilder,
+                                     int64_t handle = 0);
+
+  /**
+   * Generates an optimal trajectory.
+   *
+   * This function may take a long time to complete.
+   *
+   * @param diagnostics Enables diagnostic prints.
+   * @return Returns a holonomic trajectory on success, or a string containing a
+   *   failure reason.
+   */
+  expected<SwerveSolution, std::string> Generate(bool diagnostics = false);
+
+ private:
+  /**
+   * The swerve drivetrain.
+   */
+  SwervePath path;
+
+  /// State Variables
+  std::vector<sleipnir::Variable> x;
+  std::vector<sleipnir::Variable> y;
+  std::vector<sleipnir::Variable> thetacos;
+  std::vector<sleipnir::Variable> thetasin;
+  std::vector<sleipnir::Variable> vx;
+  std::vector<sleipnir::Variable> vy;
+  std::vector<sleipnir::Variable> omega;
+  std::vector<sleipnir::Variable> ax;
+  std::vector<sleipnir::Variable> ay;
+  std::vector<sleipnir::Variable> alpha;
+
+  /// Input Variables
+  std::vector<std::vector<sleipnir::Variable>> Fx;
+  std::vector<std::vector<sleipnir::Variable>> Fy;
+
+  /// Time Variables
+  std::vector<sleipnir::Variable> dt;
+
+  /// Discretization Constants
+  std::vector<size_t> N;
+
+  sleipnir::OptimizationProblem problem;
+  std::vector<std::function<void()>> callbacks;
+
+  void ApplyInitialGuess(const SwerveSolution& solution);
+
+  SwerveSolution ConstructSwerveSolution();
+};
+
+}  // namespace trajopt

include/trajopt/geometry/Pose2.hpp

@@ -5,6 +5,7 @@
 #include <utility>
 
 #include <sleipnir/autodiff/Variable.hpp>
+#include <sleipnir/optimization/Constraints.hpp>
 
 #include "trajopt/geometry/Rotation2.hpp"
 #include "trajopt/geometry/Translation2.hpp"
@@ -99,4 +100,16 @@ class Pose2 {
 using Pose2d = Pose2<double>;
 using Pose2v = Pose2<sleipnir::Variable>;
 
+template <typename T, typename U>
+sleipnir::EqualityConstraints operator==(const Pose2<T>& lhs,
+                                         const Pose2<U>& rhs) {
+  auto translationConstraints =
+      (lhs.Translation() == rhs.Translation()).constraints;
+  auto rotationConstraints = (lhs.Rotation() == rhs.Rotation()).constraints;
+
+  return sleipnir::EqualityConstraints{translationConstraints.insert(
+      translationConstraints.begin(), rotationConstraints.begin(),
+      rotationConstraints.end())};
+}
+
 }  // namespace trajopt

include/trajopt/geometry/Rotation2.hpp

@@ -7,6 +7,7 @@
 #include <utility>
 
 #include <sleipnir/autodiff/Variable.hpp>
+#include <sleipnir/optimization/Constraints.hpp>
 
 namespace trajopt {
 
@@ -131,4 +132,29 @@ class Rotation2 {
 using Rotation2d = Rotation2<double>;
 using Rotation2v = Rotation2<sleipnir::Variable>;
 
+template <typename T, typename U>
+sleipnir::EqualityConstraints operator==(const Rotation2<T>& lhs,
+                                         const Rotation2<U>& rhs) {
+  // Constrain angle equality on manifold.
+  //
+  // Let lhs = <cos(a), sin(a)>.  NOLINT
+  // Let rhs = <cos(b), sin(b)>.  NOLINT
+  //
+  // If the angles are equal, the angle between the unit vectors should be
+  // zero.
+  //
+  //   lhs x rhs = ||lhs|| ||rhs|| sin(angleBetween)  NOLINT
+  //         = 1 * 1 * 0
+  //         = 0
+  //
+  // NOTE: angleBetween = π rad would be another solution
+  sleipnir::VariableMatrix lhsVar{
+      {lhs.Cos() * rhs.Sin() - lhs.Sin() * rhs.Cos()},
+      // Require that lhs and rhs are unit vectors
+      {lhs.Cos() * lhs.Cos() + lhs.Sin() * lhs.Sin()},
+      {rhs.Cos() * rhs.Cos() + rhs.Sin() * rhs.Sin()}};
+  sleipnir::VariableMatrix rhsVar{{0.0}, {1.0}, {1.0}};
+  return lhsVar == rhsVar;
+}
+
 }  // namespace trajopt

include/trajopt/geometry/Translation2.hpp

@@ -7,6 +7,7 @@
 #include <utility>
 
 #include <sleipnir/autodiff/Variable.hpp>
+#include <sleipnir/optimization/Constraints.hpp>
 
 #include "trajopt/geometry/Rotation2.hpp"
 
@@ -102,6 +103,32 @@ class Translation2 {
    */
   constexpr Translation2<T> operator-() const { return {-m_x, -m_y}; }
 
+  /**
+   * Returns the translation multiplied by a scalar.
+   *
+   * @param scalar The scalar to multiply by.
+   *
+   * @return The scaled translation.
+   */
+  template <typename U>
+  constexpr auto operator*(const U& scalar) const {
+    using R = decltype(std::declval<T>() + std::declval<U>());
+    return Translation2<R>{m_x * scalar, m_y * scalar};
+  }
+
+  /**
+   * Returns the translation divided by a scalar.
+   *
+   * @param scalar The scalar to divide by.
+   *
+   * @return The scaled translation.
+   */
+  template <typename U>
+  constexpr auto operator/(const U& scalar) const {
+    using R = decltype(std::declval<T>() + std::declval<U>());
+    return Translation2<R>{m_x / scalar, m_y / scalar};
+  }
+
   /**
    * Applies a rotation to the translation in 2D space.
    *
@@ -192,6 +219,14 @@ constexpr decltype(auto) get(const trajopt::Translation2<T>& translation) {
 using Translation2d = Translation2<double>;
 using Translation2v = Translation2<sleipnir::Variable>;
 
+template <typename T, typename U>
+sleipnir::EqualityConstraints operator==(const Translation2<T>& lhs,
+                                         const Translation2<U>& rhs) {
+  sleipnir::VariableMatrix lhsVar{{lhs.X()}, {lhs.Y()}};
+  sleipnir::VariableMatrix rhsVar{{rhs.X()}, {rhs.Y()}};
+  return lhsVar == rhsVar;
+}
+
 }  // namespace trajopt
 
 namespace std {

include/trajopt/path/Path.hpp

@@ -17,46 +17,50 @@
 namespace trajopt {
 
 /**
- * A swerve path waypoint
+ * A swerve path waypoint.
  */
 struct TRAJOPT_DLLEXPORT SwerveWaypoint {
-  /// instantaneous constraints at the waypoint
+  /// Instantaneous constraints at the waypoint.
   std::vector<HolonomicConstraint> waypointConstraints;
-  /// continuous constraints along the segment
+
+  /// Continuous constraints along the segment.
   std::vector<HolonomicConstraint> segmentConstraints;
 };
 
 /**
- * A differential path waypoint
+ * A differential path waypoint.
  */
 struct TRAJOPT_DLLEXPORT DifferentialWaypoint {
-  /// instantaneous constraints at the waypoint
+  /// Instantaneous constraints at the waypoint.
   std::vector<DifferentialConstraint> waypointConstraints;
-  /// continuous constraints along the segment
+
+  /// Continuous constraints along the segment.
   std::vector<DifferentialConstraint> segmentConstraints;
 };
 
 /**
- * Swerve path
+ * Swerve path.
  */
 struct TRAJOPT_DLLEXPORT SwervePath {
-  /// waypoints along the path
+  /// Waypoints along the path.
   std::vector<SwerveWaypoint> waypoints;
-  /// drivetrain of the robot
+
+  /// Drivetrain of the robot.
   SwerveDrivetrain drivetrain;
 
   /// A vector of callbacks to be called with the intermediate SwerveSolution
-  /// and a user-specified handle at every iteration of the solver
+  /// and a user-specified handle at every iteration of the solver.
   std::vector<std::function<void(SwerveSolution&, int64_t)>> callbacks;
 };
 
 /**
- * Differential path
+ * Differential path.
  */
 struct TRAJOPT_DLLEXPORT DifferentialPath {
-  /// waypoints along the path
+  /// Waypoints along the path.
   std::vector<DifferentialWaypoint> waypoints;
-  /// drivetrain of the robot
+
+  /// Drivetrain of the robot.
   DifferentialDrivetrain drivetrain;
 };
 

include/trajopt/path/SwervePathBuilder.hpp

@@ -13,15 +13,14 @@
 #include "trajopt/obstacle/Bumpers.hpp"
 #include "trajopt/obstacle/Obstacle.hpp"
 #include "trajopt/path/Path.hpp"
-#include "trajopt/solution/Solution.hpp"
 #include "trajopt/solution/SwerveSolution.hpp"
 
 namespace trajopt {
 
 /**
  * Builds a swerve path using information about how the robot
  * must travel through a series of waypoints. This path can be converted
- * to a trajectory using OptimalTrajectoryGenerator.
+ * to a trajectory using SwerveTrajectoryGenerator.
  */
 class TRAJOPT_DLLEXPORT SwervePathBuilder {
  public:
@@ -297,7 +296,7 @@ class TRAJOPT_DLLEXPORT SwervePathBuilder {
    *
    * @return the initial guess, as a solution
    */
-  Solution CalculateInitialGuess() const;
+  SwerveSolution CalculateInitialGuess() const;
 
   /**
    * Add a callback to retrieve the state of the solver as a SwerveSolution.