Add rudimentary inverse kinematics calculator

This adds some swerve inverse kinematics that primarily serve to provide
a sanity-check and seed for more sophisticated calculations.

Change-Id: I6b214026eb53e6e47c24f4538a0b3c057f9d1751
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 637e42c..50130e8 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -469,3 +469,21 @@
         "@com_google_absl//absl/log",
     ],
 )
+
+cc_library(
+    name = "inverse_kinematics",
+    hdrs = ["inverse_kinematics.h"],
+    deps = [
+        ":simplified_dynamics",
+        "//aos/util:math",
+    ],
+)
+
+cc_test(
+    name = "inverse_kinematics_test",
+    srcs = ["inverse_kinematics_test.cc"],
+    deps = [
+        ":inverse_kinematics",
+        "//aos/testing:googletest",
+    ],
+)
diff --git a/frc971/control_loops/swerve/inverse_kinematics.h b/frc971/control_loops/swerve/inverse_kinematics.h
new file mode 100644
index 0000000..a4e43e1
--- /dev/null
+++ b/frc971/control_loops/swerve/inverse_kinematics.h
@@ -0,0 +1,86 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_INVERSE_KINEMATICS_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_INVERSE_KINEMATICS_H_
+#include "aos/util/math.h"
+#include "frc971/control_loops/swerve/simplified_dynamics.h"
+namespace frc971::control_loops::swerve {
+// Class to do straightforwards inverse kinematics of a swerve drivebase. This
+// is meant largely as a sanity-check/initializer for more sophisticated
+// methods. This calculates which directions the modules must be pointed to
+// cause them to be pointed directly along the direction of motion of the
+// drivebase. Accounting for slip angles and the such must be done as part of
+// more sophisticated inverse dynamics.
+template <typename Scalar>
+class InverseKinematics {
+ public:
+  using ModuleParams = SimplifiedDynamics<Scalar>::ModuleParams;
+  using Parameters = SimplifiedDynamics<Scalar>::Parameters;
+  using States = SimplifiedDynamics<Scalar>::States;
+  using State = SimplifiedDynamics<Scalar>::template VelocityState<Scalar>;
+  InverseKinematics(const Parameters &params) : params_(params) {}
+
+  // Uses kVx, kVy, kTheta, and kOmega from the input goal state for the
+  // absolute kinematics. Also uses the specified theta values to bias theta
+  // output values towards the current state (i.e., if the module 0 theta is
+  // currently 0 and we are asked to drive straight backwards, this will prefer
+  // a theta of zero rather than a theta of pi).
+  State Solve(const State &goal) {
+    State result = goal;
+    for (size_t module_index = 0; module_index < params_.modules.size();
+         ++module_index) {
+      SolveModule(goal, params_.modules[module_index].position,
+                  &result(States::kThetas0 + 2 * module_index),
+                  &result(States::kOmegas0 + 2 * module_index));
+    }
+    return result;
+  }
+
+  void SolveModule(const State &goal,
+                   const Eigen::Matrix<Scalar, 2, 1> &module_position,
+                   Scalar *module_theta, Scalar *module_omega) {
+    const Scalar vx = goal(States::kVx);
+    const Scalar vy = goal(States::kVy);
+    const Scalar omega = goal(States::kOmega);
+    // module_velocity_in_robot_frame = R(-theta) * robot_vel +
+    //     omega.cross(module_position);
+    // module_vel_x = (cos(-theta) * vx - sin(-theta) * vy) - omega * module_y
+    // module_vel_y = (sin(-theta) * vx + cos(-theta) * vy) + omega * module_x
+    // module_theta = atan2(module_vel_y, module_vel_x)
+    // module_omega = datan2(module_vel_y, module_vel_x) / dt
+    // datan2(y, x) / dt = (x * dy/dt - y * dx / dt) / (x^2 + y^2)
+    // robot accelerations are assumed to be zero.
+    // dmodule_vel_x / dt = (sin(-theta) * vx + cos(-theta) * vy) * omega
+    // dmodule_vel_y / dt = (-cos(-theta) * vx + sin(-theta) * vy) * omega
+    const Scalar ctheta = std::cos(-goal(States::kTheta));
+    const Scalar stheta = std::sin(-goal(States::kTheta));
+    const Scalar module_vel_x =
+        (ctheta * vx - stheta * vy) - omega * module_position.y();
+    const Scalar module_vel_y =
+        (stheta * vx + ctheta * vy) + omega * module_position.x();
+    const Scalar nominal_module_theta = atan2(module_vel_y, module_vel_x);
+    // If the current module theta is more than 90 deg from the desired theta,
+    // flip the desired theta by 180 deg.
+    if (std::abs(aos::math::DiffAngle(nominal_module_theta, *module_theta)) >
+        M_PI_2) {
+      *module_theta = aos::math::NormalizeAngle(nominal_module_theta + M_PI);
+    } else {
+      *module_theta = nominal_module_theta;
+    }
+    const Scalar module_accel_x = (stheta * vx + ctheta * vy) * omega;
+    const Scalar module_accel_y = (-ctheta * vx + stheta * vy) * omega;
+    const Scalar module_vel_norm_squared =
+        (module_vel_x * module_vel_x + module_vel_y * module_vel_y);
+    if (module_vel_norm_squared < 1e-5) {
+      // Prevent poor conditioning of module velocities at near-zero speeds.
+      *module_omega = 0.0;
+    } else {
+      *module_omega =
+          (module_vel_x * module_accel_y - module_vel_y * module_accel_x) /
+          module_vel_norm_squared;
+    }
+  }
+
+ private:
+  Parameters params_;
+};
+}  // namespace frc971::control_loops::swerve
+#endif  // FRC971_CONTROL_LOOPS_SWERVE_INVERSE_KINEMATICS_H_
diff --git a/frc971/control_loops/swerve/inverse_kinematics_test.cc b/frc971/control_loops/swerve/inverse_kinematics_test.cc
new file mode 100644
index 0000000..407d540
--- /dev/null
+++ b/frc971/control_loops/swerve/inverse_kinematics_test.cc
@@ -0,0 +1,150 @@
+#include "frc971/control_loops/swerve/inverse_kinematics.h"
+
+#include "gtest/gtest.h"
+
+namespace frc971::control_loops::swerve::testing {
+class InverseKinematicsTest : public ::testing::Test {
+ protected:
+  typedef double Scalar;
+  using State = InverseKinematics<Scalar>::State;
+  using States = InverseKinematics<Scalar>::States;
+  using ModuleParams = InverseKinematics<Scalar>::ModuleParams;
+  using Parameters = InverseKinematics<Scalar>::Parameters;
+  static ModuleParams MakeModule(const Eigen::Matrix<Scalar, 2, 1> &position) {
+    return ModuleParams{.position = position,
+                        .slip_angle_coefficient = 0.0,
+                        .slip_angle_alignment_coefficient = 0.0,
+                        .steer_motor = KrakenFOC(),
+                        .drive_motor = KrakenFOC(),
+                        .steer_ratio = 1.0,
+                        .drive_ratio = 1.0};
+  }
+  static Parameters MakeParams() {
+    return {.mass = 1.0,
+            .moment_of_inertia = 1.0,
+            .modules = {
+                MakeModule({1.0, 1.0}),
+                MakeModule({-1.0, 1.0}),
+                MakeModule({-1.0, -1.0}),
+                MakeModule({1.0, -1.0}),
+            }};
+  }
+
+  InverseKinematicsTest() : inverse_kinematics_(MakeParams()) {}
+
+  struct Goal {
+    Scalar vx;
+    Scalar vy;
+    Scalar omega;
+    Scalar theta;
+  };
+
+  void CheckState(
+      const Goal &goal, const std::array<Scalar, 4> &expected_thetas,
+      const std::optional<Eigen::Vector4d> &expected_omegas = std::nullopt) {
+    State goal_state = State::Zero();
+    goal_state(States::kVx) = goal.vx;
+    goal_state(States::kVy) = goal.vy;
+    goal_state(States::kOmega) = goal.omega;
+    goal_state(States::kTheta) = goal.theta;
+    SCOPED_TRACE(goal_state.bottomRows<4>().transpose());
+    const State nominal_state = inverse_kinematics_.Solve(goal_state);
+    // Now, calculate the numerical derivative of the state and validate that it
+    // matches expectations.
+    const Scalar kDt = 1e-5;
+    const Scalar dtheta = kDt * goal.omega;
+    goal_state(States::kTheta) += dtheta / 2.0;
+    const State state_eps_pos = inverse_kinematics_.Solve(goal_state);
+    goal_state(States::kTheta) -= dtheta;
+    const State state_eps_neg = inverse_kinematics_.Solve(goal_state);
+    const State state_derivative = (state_eps_pos - state_eps_neg) / kDt;
+    for (size_t module_index = 0; module_index < 4; ++module_index) {
+      SCOPED_TRACE(module_index);
+      const int omega_idx = States::kOmegas0 + 2 * module_index;
+      const int theta_idx = States::kThetas0 + 2 * module_index;
+      EXPECT_NEAR(nominal_state(omega_idx), state_derivative(theta_idx), 1e-10);
+      EXPECT_NEAR(nominal_state(theta_idx), expected_thetas[module_index],
+                  1e-10);
+      if (expected_omegas.has_value()) {
+        EXPECT_NEAR(nominal_state(omega_idx),
+                    expected_omegas.value()(module_index), 1e-10);
+      }
+    }
+  }
+
+  InverseKinematics<Scalar> inverse_kinematics_;
+};
+
+// Tests that if we are driving straight with no yaw that we get sane
+// kinematics.
+TEST_F(InverseKinematicsTest, StraightDrivingNoYaw) {
+  // Sanity-check zero-speed operation.
+  CheckState({.vx = 0.0, .vy = 0.0, .omega = 0.0, .theta = 0.0},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = 0.0},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+  // Reverse should prefer to bias the modules towards [-pi/2, pi/2] due to
+  // hysteresis from starting the modules at thetas of 0.
+  CheckState({.vx = -1.0, .vy = 0.0, .omega = 0.0, .theta = 0.0},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 0.0, .vy = 1.0, .omega = 0.0, .theta = 0.0},
+             {M_PI_2, M_PI_2, M_PI_2, M_PI_2}, Eigen::Vector4d::Zero());
+  // For module hysteresis, this is a corner case where we are exactly 90 deg
+  // from the current value; the exact result is unimportant.
+  CheckState({.vx = 0.0, .vy = -1.0, .omega = 0.0, .theta = 0.0},
+             {-M_PI_2, -M_PI_2, -M_PI_2, -M_PI_2}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 1.0, .vy = 1.0, .omega = 0.0, .theta = 0.0},
+             {M_PI_4, M_PI_4, M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+  // Reverse should prefer to bias the modules towards [-pi/2, pi/2] due to
+  // hysteresis from starting the modules at thetas of 0.
+  CheckState({.vx = -1.0, .vy = -1.0, .omega = 0.0, .theta = 0.0},
+             {M_PI_4, M_PI_4, M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+}
+
+// Tests that if we are driving straight with non-zero yaw that we get sane
+// kinematics.
+TEST_F(InverseKinematicsTest, StraightDrivingYawed) {
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = 0.1},
+             {-0.1, -0.1, -0.1, -0.1}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = M_PI_2},
+             {-M_PI_2, -M_PI_2, -M_PI_2, -M_PI_2}, Eigen::Vector4d::Zero());
+  CheckState({.vx = 1.0, .vy = 0.0, .omega = 0.0, .theta = M_PI},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+
+  CheckState({.vx = 0.0, .vy = 1.0, .omega = 0.0, .theta = M_PI_2},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+  // Reverse should prefer to bias the modules towards [-pi/2, pi/2] due to
+  // hysteresis from starting the modules at thetas of 0.
+  CheckState({.vx = 0.0, .vy = -1.0, .omega = 0.0, .theta = M_PI_2},
+             {0.0, 0.0, 0.0, 0.0}, Eigen::Vector4d::Zero());
+}
+
+// Tests that we can spin in place.
+TEST_F(InverseKinematicsTest, SpinInPlace) {
+  CheckState({.vx = 0.0, .vy = 0.0, .omega = 1.0, .theta = 0.0},
+             {-M_PI_4, M_PI_4, -M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+  // And changing the current theta should not matter.
+  CheckState({.vx = 0.0, .vy = 0.0, .omega = 1.0, .theta = 1.0},
+             {-M_PI_4, M_PI_4, -M_PI_4, M_PI_4}, Eigen::Vector4d::Zero());
+}
+
+// Tests that if we are spinning while moving that we correctly calculate module
+// yaw rates.
+TEST_F(InverseKinematicsTest, SpinWhileMoving) {
+  // Set up a situation where we are driving straight forwards, with a
+  // yaw rate of 1 rad / sec.
+  // The modules are all at radii of sqrt(2), so the contribution from both
+  // the translational and rotational velocities should be equal; in this case
+  // each module will have an angle that is an equal combination of straight
+  // forwards (0 deg) and some 45 deg offset.
+  CheckState({.vx = std::sqrt(static_cast<Scalar>(2)),
+              .vy = 0.0,
+              .omega = 1.0,
+              .theta = 0.0},
+             {3 * M_PI_4 / 2, -3 * M_PI_4 / 2, -M_PI_4 / 2, M_PI_4 / 2});
+}
+}  // namespace frc971::control_loops::swerve::testing