Add FRC 2024 shooter example (#583)

examples/FRC2022Shooter/main.py

@@ -3,8 +3,8 @@
 """
 FRC 2022 shooter trajectory optimization.
 
-This program finds the optimal initial launch velocity and launch angle for the
-2022 FRC game's target.
+This program finds the initial velocity, pitch, and yaw for a game piece to hit
+the 2022 FRC game's target that minimizes time-to-target.
 """
 
 import math
@@ -21,21 +21,44 @@
 target_wrt_field = np.array(
     [[field_length / 2.0], [field_width / 2.0], [2.64], [0.0], [0.0], [0.0]]
 )
-target_radius = 0.61
+target_radius = 0.61  # m
 g = 9.806  # m/s²
 
 
 def lerp(a, b, t):
     return a + t * (b - a)
 
 
+def f(x):
+    # x' = x'
+    # y' = y'
+    # z' = z'
+    # x" = −a_D(v_x)
+    # y" = −a_D(v_y)
+    # z" = −g − a_D(v_z)
+    #
+    # where a_D(v) = ½ρv² C_D A / m
+    rho = 1.204  # kg/m³
+    C_D = 0.5
+    A = math.pi * 0.3
+    m = 2.0  # kg
+    a_D = lambda v: 0.5 * rho * v**2 * C_D * A / m
+
+    v_x = x[3, 0]
+    v_y = x[4, 0]
+    v_z = x[5, 0]
+    return VariableMatrix(
+        [[v_x], [v_y], [v_z], [-a_D(v_x)], [-a_D(v_y)], [-g - a_D(v_z)]]
+    )
+
+
 def main():
     # Robot initial state
     robot_wrt_field = np.array(
         [[field_length / 4.0], [field_width / 4.0], [0.0], [1.524], [-1.524], [0.0]]
     )
 
-    max_launch_velocity = 10.0  # m/s
+    max_initial_velocity = 10.0  # m/s
 
     shooter_wrt_robot = np.array([[0.0], [0.0], [1.2], [0.0], [0.0], [0.0]])
     shooter_wrt_field = robot_wrt_field + shooter_wrt_robot
@@ -75,59 +98,31 @@ def main():
         p_y[k].set_value(lerp(shooter_wrt_field[1, 0], target_wrt_field[1, 0], k / N))
         p_z[k].set_value(lerp(shooter_wrt_field[2, 0], target_wrt_field[2, 0], k / N))
 
-    # Velocity initial guess is max launch velocity toward goal
+    # Velocity initial guess is max initial velocity toward target
     uvec_shooter_to_target = target_wrt_field[:3, :] - shooter_wrt_field[:3, :]
     uvec_shooter_to_target /= norm(uvec_shooter_to_target)
     for k in range(N):
         v[:, k].set_value(
-            robot_wrt_field[3:, :] + max_launch_velocity * uvec_shooter_to_target
+            robot_wrt_field[3:, :] + max_initial_velocity * uvec_shooter_to_target
         )
 
     # Shooter initial position
     problem.subject_to(p[:, 0] == shooter_wrt_field[:3, 0])
 
-    # Require initial launch velocity is below max
+    # Require initial velocity is below max
     #
     #   √{v_x² + v_y² + v_z²) ≤ vₘₐₓ
     #   v_x² + v_y² + v_z² ≤ vₘₐₓ²
     problem.subject_to(
         (v_x[0] - robot_wrt_field[3, 0]) ** 2
         + (v_y[0] - robot_wrt_field[4, 0]) ** 2
         + (v_z[0] - robot_wrt_field[5, 0]) ** 2
-        <= max_launch_velocity**2
+        <= max_initial_velocity**2
     )
 
-    # Require final position is in center of target circle
-    problem.subject_to(p[:, -1] == target_wrt_field[:3, :])
-
-    # Require the final velocity is down
-    problem.subject_to(v_z[-1] < 0.0)
-
-    def f(x):
-        # x' = x'
-        # y' = y'
-        # z' = z'
-        # x" = −a_D(v_x)
-        # y" = −a_D(v_y)
-        # z" = −g − a_D(v_z)
-        #
-        # where a_D(v) = ½ρv² C_D A / m
-        rho = 1.204  # kg/m³
-        C_D = 0.5
-        A = math.pi * 0.3
-        m = 2.0  # kg
-        a_D = lambda v: 0.5 * rho * v**2 * C_D * A / m
-
-        v_x = x[3, 0]
-        v_y = x[4, 0]
-        v_z = x[5, 0]
-        return VariableMatrix(
-            [[v_x], [v_y], [v_z], [-a_D(v_x)], [-a_D(v_y)], [-g - a_D(v_z)]]
-        )
-
     # Dynamics constraints - RK4 integration
+    h = dt
     for k in range(N - 1):
-        h = dt
         x_k = X[:, k]
         x_k1 = X[:, k + 1]
 
@@ -137,16 +132,22 @@ def f(x):
         k4 = f(x_k + h * k3)
         problem.subject_to(x_k1 == x_k + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4))
 
-    # Minimize time to goal
+    # Require final position is in center of target circle
+    problem.subject_to(p[:, -1] == target_wrt_field[:3, :])
+
+    # Require the final velocity is down
+    problem.subject_to(v_z[-1] < 0.0)
+
+    # Minimize time-to-target
     problem.minimize(T)
 
     problem.solve(diagnostics=True)
 
     # Initial velocity vector
     v0 = X[3:, 0].value() - robot_wrt_field[3:, :]
 
-    launch_velocity = norm(v0)
-    print(f"Launch velocity = {launch_velocity:.03f} m/s")
+    velocity = norm(v0)
+    print(f"Velocity = {velocity:.03f} m/s")
 
     pitch = math.atan2(v0[2, 0], math.hypot(v0[0, 0], v0[1, 0]))
     print(f"Pitch = {np.rad2deg(pitch):.03f}°")

examples/FRC2022Shooter/src/Main.cpp

@@ -7,10 +7,10 @@
 #include <Eigen/Core>
 #include <sleipnir/optimization/OptimizationProblem.hpp>
 
-// FRC 2022 shooter trajectory optimization
+// FRC 2022 shooter trajectory optimization.
 //
-// This program finds the optimal initial launch velocity and launch angle for
-// the 2022 FRC game's target.
+// This program finds the initial velocity, pitch, and yaw for a game piece to
+// hit the 2022 FRC game's target that minimizes time-to-target.
 
 namespace slp = sleipnir;
 
@@ -21,7 +21,31 @@ constexpr double field_width = 8.2296;    // 27 ft -> m
 constexpr double field_length = 16.4592;  // 54 ft -> m
 constexpr Vector6d target_wrt_field{
     {field_length / 2.0}, {field_width / 2.0}, {2.64}, {0.0}, {0.0}, {0.0}};
-constexpr double g = 9.806;  // m/s²
+[[maybe_unused]]
+constexpr double target_radius = 0.61;  // m
+constexpr double g = 9.806;             // m/s²
+
+slp::VariableMatrix f(const slp::VariableMatrix& x) {
+  // x' = x'
+  // y' = y'
+  // z' = z'
+  // x" = −a_D(v_x)
+  // y" = −a_D(v_y)
+  // z" = −g − a_D(v_z)
+  //
+  // where a_D(v) = ½ρv² C_D A / m
+  constexpr double rho = 1.204;  // kg/m³
+  constexpr double C_D = 0.5;
+  constexpr double A = std::numbers::pi * 0.3;
+  constexpr double m = 2.0;  // kg
+  auto a_D = [](auto v) { return 0.5 * rho * v * v * C_D * A / m; };
+
+  auto v_x = x(3, 0);
+  auto v_y = x(4, 0);
+  auto v_z = x(5, 0);
+  return slp::VariableMatrix{{v_x},       {v_y},       {v_z},
+                             {-a_D(v_x)}, {-a_D(v_y)}, {-g - a_D(v_z)}};
+}
 
 #ifndef RUNNING_TESTS
 int main() {
@@ -33,7 +57,7 @@ int main() {
                                      {-1.524},
                                      {0.0}};
 
-  constexpr double max_launch_velocity = 10.0;
+  constexpr double max_initial_velocity = 10.0;  // m/s
 
   Vector6d shooter_wrt_robot{{0.0}, {0.0}, {1.2}, {0.0}, {0.0}, {0.0}};
   Vector6d shooter_wrt_field = robot_wrt_field + shooter_wrt_robot;
@@ -78,58 +102,30 @@ int main() {
                               static_cast<double>(k) / N));
   }
 
-  // Velocity initial guess is max launch velocity toward goal
+  // Velocity initial guess is max initial velocity toward target
   Vector3d uvec_shooter_to_target =
-      (target_wrt_field.block(0, 0, 3, 1) - shooter_wrt_field.block(0, 0, 3, 1))
+      (target_wrt_field.segment(0, 3) - shooter_wrt_field.segment(0, 3))
           .normalized();
   for (int k = 0; k < N; ++k) {
-    v.Col(k).SetValue(robot_wrt_field.block(3, 0, 3, 1) +
-                      max_launch_velocity * uvec_shooter_to_target);
+    v.Col(k).SetValue(robot_wrt_field.segment(3, 3) +
+                      max_initial_velocity * uvec_shooter_to_target);
   }
 
   // Shooter initial position
   problem.SubjectTo(p.Col(0) == shooter_wrt_field.block(0, 0, 3, 1));
 
-  // Require initial launch velocity is below max
+  // Require initial velocity is below max
   //
   //   √{v_x² + v_y² + v_z²) ≤ vₘₐₓ
   //   v_x² + v_y² + v_z² ≤ vₘₐₓ²
   problem.SubjectTo(slp::pow(v_x(0) - robot_wrt_field(3), 2) +
                         slp::pow(v_y(0) - robot_wrt_field(4), 2) +
                         slp::pow(v_z(0) - robot_wrt_field(5), 2) <=
-                    max_launch_velocity * max_launch_velocity);
-
-  // Require final position is in center of target circle
-  problem.SubjectTo(p.Col(N - 1) == target_wrt_field.block(0, 0, 3, 1));
-
-  // Require the final velocity is down
-  problem.SubjectTo(v_z(N - 1) < 0.0);
-
-  auto f = [](slp::VariableMatrix x) {
-    // x' = x'
-    // y' = y'
-    // z' = z'
-    // x" = −a_D(v_x)
-    // y" = −a_D(v_y)
-    // z" = −g − a_D(v_z)
-    //
-    // where a_D(v) = ½ρv² C_D A / m
-    constexpr double rho = 1.204;  // kg/m³
-    constexpr double C_D = 0.5;
-    constexpr double A = std::numbers::pi * 0.3;
-    constexpr double m = 2.0;  // kg
-    auto a_D = [](auto v) { return 0.5 * rho * v * v * C_D * A / m; };
-
-    auto v_x = x(3, 0);
-    auto v_y = x(4, 0);
-    auto v_z = x(5, 0);
-    return slp::VariableMatrix{{v_x},       {v_y},       {v_z},
-                               {-a_D(v_x)}, {-a_D(v_y)}, {-g - a_D(v_z)}};
-  };
+                    max_initial_velocity * max_initial_velocity);
 
   // Dynamics constraints - RK4 integration
+  auto h = dt;
   for (int k = 0; k < N - 1; ++k) {
-    auto h = dt;
     auto x_k = X.Col(k);
     auto x_k1 = X.Col(k + 1);
 
@@ -140,17 +136,23 @@ int main() {
     problem.SubjectTo(x_k1 == x_k + h / 6 * (k1 + 2 * k2 + 2 * k3 + k4));
   }
 
-  // Minimize time to goal
+  // Require final position is in center of target circle
+  problem.SubjectTo(p.Col(N - 1) == target_wrt_field.block(0, 0, 3, 1));
+
+  // Require the final velocity is down
+  problem.SubjectTo(v_z(N - 1) < 0.0);
+
+  // Minimize time-to-target
   problem.Minimize(T);
 
   problem.Solve({.diagnostics = true});
 
   // Initial velocity vector
   Eigen::Vector3d v0 =
-      X.Block(3, 0, 3, 1).Value() - robot_wrt_field.block(3, 0, 3, 1);
+      X.Block(3, 0, 3, 1).Value() - robot_wrt_field.segment(3, 3);
 
-  double launch_velocity = v0.norm();
-  std::println("Launch velocity = {:.03} ms", launch_velocity);
+  double velocity = v0.norm();
+  std::println("Velocity = {:.03} ms", velocity);
 
   double pitch = std::atan2(v0(2), std::hypot(v0(0), v0(1)));
   std::println("Pitch = {:.03}°", pitch * 180.0 / std::numbers::pi);