Merge "Add convenient auto-diff Jacobian calculator" into main
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 7a8c1ea..224b8e8 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -398,3 +398,21 @@
"@com_google_absl//absl/log:check",
],
)
+
+cc_library(
+ name = "auto_diff_jacobian",
+ hdrs = ["auto_diff_jacobian.h"],
+ target_compatible_with = ["@platforms//os:linux"],
+ deps = [
+ "@com_google_ceres_solver//:ceres",
+ ],
+)
+
+cc_test(
+ name = "auto_diff_jacobian_test",
+ srcs = ["auto_diff_jacobian_test.cc"],
+ deps = [
+ ":auto_diff_jacobian",
+ "//aos/testing:googletest",
+ ],
+)
diff --git a/frc971/control_loops/swerve/auto_diff_jacobian.h b/frc971/control_loops/swerve/auto_diff_jacobian.h
new file mode 100644
index 0000000..2e7947a
--- /dev/null
+++ b/frc971/control_loops/swerve/auto_diff_jacobian.h
@@ -0,0 +1,62 @@
+#ifndef FRC971_CONTROL_LOOPS_SWERVE_AUTO_DIFF_JACOBIAN_H_
+#define FRC971_CONTROL_LOOPS_SWERVE_AUTO_DIFF_JACOBIAN_H_
+#include "include/ceres/tiny_solver.h"
+#include "include/ceres/tiny_solver_autodiff_function.h"
+
+namespace frc971::control_loops::swerve {
+// Class to conveniently scope a function that makes use of Ceres'
+// autodifferentiation methods to calculate the jacobian of the provided method.
+// Template parameters:
+// Scalar: scalar type to use (typically double or float; this is used for
+// allowing you to control what precision you use; you cannot use this
+// with ceres Jets because this has to use Jets internally).
+// Function: The type of the function itself. A Function f must be callable as
+// Eigen::Matrix<Scalar, kNumOutputs, 1> = f(Eigen::Matrix<Scalar,
+// kNumInputs, 1>{});
+template <typename Scalar, typename Function, size_t kNumInputs,
+ size_t kNumOutputs>
+class AutoDiffJacobian {
+ public:
+ // Calculates the jacobian of the provided method, function, at the provided
+ // input X.
+ static Eigen::Matrix<Scalar, kNumOutputs, kNumInputs> Jacobian(
+ const Function &function, const Eigen::Matrix<Scalar, kNumInputs, 1> &X) {
+ AutoDiffCeresFunctor ceres_functor(function);
+ TinySolverFunctor tiny_solver(ceres_functor);
+ // residual is unused, it's just a place to store the evaluated function at
+ // the current state/input.
+ Eigen::Matrix<Scalar, kNumOutputs, 1> residual;
+ Eigen::Matrix<Scalar, kNumOutputs, kNumInputs> jacobian;
+ tiny_solver(X.data(), residual.data(), jacobian.data());
+ return jacobian;
+ }
+
+ private:
+ // Borrow the TinySolver's auto-differentiation execution for use here to
+ // calculate the linearized dynamics. We aren't actually doing any solving,
+ // just letting it do the jacobian calculation for us.
+ // As such, construct a "residual" function whose residuals are just the
+ // derivative of the state and whose parameters are the stacked state + input.
+ class AutoDiffCeresFunctor {
+ public:
+ AutoDiffCeresFunctor(const Function &function) : function_(function) {}
+ template <typename ScalarT>
+ bool operator()(const ScalarT *const parameters,
+ ScalarT *const residuals) const {
+ const Eigen::Map<const Eigen::Matrix<ScalarT, kNumInputs, 1>>
+ eigen_parameters(parameters);
+ Eigen::Map<Eigen::Matrix<ScalarT, kNumOutputs, 1>> eigen_residuals(
+ residuals);
+ eigen_residuals = function_(eigen_parameters);
+ return true;
+ }
+
+ private:
+ const Function &function_;
+ };
+ typedef ceres::TinySolverAutoDiffFunction<AutoDiffCeresFunctor, kNumOutputs,
+ kNumInputs, Scalar>
+ TinySolverFunctor;
+};
+} // namespace frc971::control_loops::swerve
+#endif // FRC971_CONTROL_LOOPS_SWERVE_AUTO_DIFF_JACOBIAN_H_
diff --git a/frc971/control_loops/swerve/auto_diff_jacobian_test.cc b/frc971/control_loops/swerve/auto_diff_jacobian_test.cc
new file mode 100644
index 0000000..d04726c
--- /dev/null
+++ b/frc971/control_loops/swerve/auto_diff_jacobian_test.cc
@@ -0,0 +1,23 @@
+#include "frc971/control_loops/swerve/auto_diff_jacobian.h"
+
+#include <functional>
+
+#include "absl/log/log.h"
+#include "gtest/gtest.h"
+
+namespace frc971::control_loops::swerve::testing {
+struct TestFunction {
+ template <typename Scalar>
+ Eigen::Matrix<Scalar, 3, 1> operator()(
+ const Eigen::Map<const Eigen::Matrix<Scalar, 2, 1>> X) const {
+ return Eigen::Matrix<double, 3, 2>{{1, 2}, {3, 4}, {5, 6}} * X;
+ }
+};
+
+TEST(AutoDiffJacobianTest, EvaluatesJacobian) {
+ EXPECT_EQ((AutoDiffJacobian<double, TestFunction, 2, 3>::Jacobian(
+ TestFunction{}, Eigen::Vector2d::Zero())),
+ (Eigen::Matrix<double, 3, 2>{{1, 2}, {3, 4}, {5, 6}}));
+}
+
+} // namespace frc971::control_loops::swerve::testing