feat: autodiff does more derivatives (also more comprehensive tests)

CMakeLists.txt

@@ -90,18 +90,17 @@ endif()
 
 if(BUILD_EXAMPLES)
     file(GLOB example_subdirs "${CMAKE_CURRENT_SOURCE_DIR}/examples/*")
-    foreach(example ${example_subdirs})
-        message(${example})
-        if(IS_DIRECTORY ${example})
+    foreach(example example_subdirs)
+        if(IS_DIRECTORY example)
             cmake_path(GET example FILENAME example_name)
             file(GLOB_RECURSE example_sources "${example}/*.cpp")
-            add_executable(${example_name} ${example_sources})
+            add_executable(example_name example_sources)
             target_compile_options(
-                ${example_name}
+                example_name
                 PRIVATE -Wall -Wextra -Wpedantic -Werror
             )
-            target_include_directories(${example_name} PRIVATE "${example}")
-            target_link_libraries(${example_name} PRIVATE suboptimal)
+            target_include_directories(example_name PRIVATE "${example}")
+            target_link_libraries(example_name PRIVATE suboptimal)
         endif()
     endforeach()
 endif()

include/suboptimal/autodiff/Variable.h

@@ -25,7 +25,6 @@ using Matrix2v = Eigen::Matrix2<Variable>;
 using Matrix3v = Eigen::Matrix3<Variable>;
 using Matrix4v = Eigen::Matrix4<Variable>;
 
-
 /**
  * An autodiff variable. Essentially just a nicer wrapper around Expression
  */
@@ -203,7 +202,6 @@ Variable tan(const Variable& x);
 Variable asin(const Variable& x);
 Variable acos(const Variable& x);
 Variable atan(const Variable& x);
-Variable atan2(const Variable& y, const Variable& x);
 
 template <typename Y, typename X>
   requires VariableLike<Y> && VariableLike<X> && (std::same_as<Y, Variable> || std::same_as<X, Variable>)

include/suboptimal/autodiff/derivatives.h

@@ -1,15 +1,43 @@
+// Copyright (c) 2024 Alvin Zhang.
+
 #pragma once
 
 #include <Eigen/Core>
 
 #include "suboptimal/autodiff/Variable.h"
 
 namespace suboptimal {
+/**
+ * Computes the gradient of a variable with respect to a set of variables
+ * @param var the variable to compute the gradient of
+ * @param wrt the variables to compute the gradient with respect to
+ * @return a vector of variables representing the gradient of var with respect to wrt
+ */
 VectorXv gradient(const Variable& var, const Eigen::Ref<const VectorXv>& wrt);
 
+/**
+ * Computes the derivative of a variable with respect to another variable
+ * @param var the variable to compute the derivative of
+ * @param wrt the variable to compute the derivative with respect to
+ * @return a variable representing the derivative of var with respect to wrt
+ */
 Variable derivative(const Variable& var, const Variable& wrt);
 
+/**
+ * Computes the Jacobian of a vector of variables with respect to another vector of variables
+ * @param vars the vector variables to compute the Jacobian of
+ * @param wrt the vector of variables to compute the Jacobian with respect to
+ * @return an n x m matrix of variables representing the Jacobian of vars with respect to wrt, where n is the length of
+ * vars and m is the length of wrt
+ */
 MatrixXv jacobian(const Eigen::Ref<const VectorXv>& vars, const Eigen::Ref<const VectorXv>& wrt);
 
+/**
+ * Computes the Hessian of a variable with respect to a set of variables
+ * @param var the variable to compute the Hessian of
+ * @param wrt the variables to compute the Hessian with respect to
+ * @return an n x n matrix of variables representing the Hessian of var with respect to wrt, where n is the length of
+ * wrt
+ */
 MatrixXv hessian(const Variable& var, const Eigen::Ref<const VectorXv>& wrt);
-}  // namespace suboptimal
+}  // namespace suboptimal

src/autodiff/Expression.cpp

@@ -13,7 +13,7 @@ namespace suboptimal {
 Expression::Expression(const double value, const ExpressionType type) : value{value}, type{type} {}
 
 Expression::Expression(const ExpressionType type, const ValueFunc value_func, const AdjointValueFunc adjoint_value_func,
-                       const AdjointExprFunc adjoint_expr_func, const ExpressionPtr arg) // NOLINT
+                       const AdjointExprFunc adjoint_expr_func, const ExpressionPtr arg)  // NOLINT
     : value{value_func(arg->value, 0.0)},
       lhs{arg},
       value_func{value_func},
@@ -214,9 +214,9 @@ ExpressionPtr abs(const ExpressionPtr& x) {
 
   return std::make_shared<Expression>(
       ExpressionType::Nonlinear, [](const double val, double) { return std::abs(val); },
-      [](const double val, double, const double parent_adjoint) { return val < 0 ? -parent_adjoint : parent_adjoint; },
+      [](const double val, double, const double parent_adjoint) { return parent_adjoint * val / std::abs(val); },
       [](const ExpressionPtr& expr, const ExpressionPtr&, const ExpressionPtr& parent_adjoint) {
-        return expr->value < 0 ? -parent_adjoint : parent_adjoint;
+        return parent_adjoint * expr / suboptimal::abs(expr);
       },
       x);
 }

src/autodiff/derivatives.cpp

@@ -1,5 +1,9 @@
+// Copyright (c) 2024 Alvin Zhang.
+
 #include "suboptimal/autodiff/derivatives.h"
 
+#include <memory>
+
 #include <Eigen/Core>
 
 namespace suboptimal {
@@ -52,4 +56,16 @@ VectorXv gradient(const Variable& var, const Eigen::Ref<const VectorXv>& wrt) {
 Variable derivative(const Variable& var, const Variable& wrt) {
   return gradient(var, VectorXv{{wrt}})(0);
 }
+
+MatrixXv jacobian(const Eigen::Ref<const VectorXv>& vars, const Eigen::Ref<const VectorXv>& wrt) {
+  MatrixXv jacobian{vars.size(), wrt.size()};
+  for (int i = 0; i < vars.size(); i++) {
+    jacobian.row(i) = gradient(vars(i), wrt).transpose();
+  }
+  return jacobian;
+}
+
+MatrixXv hessian(const Variable& var, const Eigen::Ref<const VectorXv>& wrt) {
+  return jacobian(gradient(var, wrt), wrt);
+}
 }  // namespace suboptimal

test/autodiff/gradient_test.cpp

@@ -0,0 +1,175 @@
+// Copyright (c) 2024 Alvin Zhang.
+
+#define CATCH_CONFIG_FAST_COMPILE
+
+#include <suboptimal/autodiff/Variable.h>
+#include <suboptimal/autodiff/derivatives.h>
+
+#include <cmath>
+
+#include <Eigen/Core>
+#include <catch2/catch_test_macros.hpp>
+#include <catch2/generators/catch_generators.hpp>
+#include <catch2/generators/catch_generators_adapters.hpp>
+#include <catch2/generators/catch_generators_random.hpp>
+#include <catch2/matchers/catch_matchers.hpp>
+#include <catch2/matchers/catch_matchers_floating_point.hpp>
+
+using namespace suboptimal;
+
+TEST_CASE("Autodiff - Basic gradient", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  const Variable f = (x * x - 2 * x * y + y * y) / (5 * x);
+  const Vector2v grad = gradient(f, Vector2v{x, y});
+
+  const double x_val = GENERATE(take(10, random(-100.0, 100.0)));
+  const double y_val = GENERATE(take(10, random(-100.0, 100.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+  CHECK(suboptimal::getValues(grad).isApprox(Eigen::Vector2d{(x_val * x_val - y_val * y_val) / (5 * x_val * x_val),  //
+                                                             (2 * y_val - 2 * x_val) / (5 * x_val)}));
+}
+
+TEST_CASE("Autodiff - Gradient of abs", "[autodiff]") {
+  Variable x{};
+  const Variable f = suboptimal::abs(x);
+  const auto grad = gradient(f, VectorXv{{x}});
+
+  const double x_val = GENERATE(take(10, random(-100.0, 100.0)));
+  x.setValue(x_val);
+  CHECK_THAT(suboptimal::getValues(grad)(0), Catch::Matchers::WithinAbs(x_val / std::abs(x_val), 1e-9));
+}
+
+TEST_CASE("Autodiff - Gradient of sqrt", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  Variable z{};
+  const Variable f = suboptimal::sqrt(x + 2 * y - z);
+  const Vector3v grad = gradient(f, Vector3v{x, y, z});
+
+  const double x_val = GENERATE(take(5, random(0.0, 100.0)));
+  const double y_val = GENERATE(take(5, random(0.0, 100.0)));
+  const double z_val = GENERATE(take(5, random(-100.0, 0.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+  z.setValue(z_val);
+  CHECK(suboptimal::getValues(grad).isApprox(Eigen::Vector3d{0.5 / std::sqrt(x_val + 2 * y_val - z_val),
+                                                             1 / std::sqrt(x_val + 2 * y_val - z_val),
+                                                             -0.5 / std::sqrt(x_val + 2 * y_val - z_val)}));
+}
+
+TEST_CASE("Autodiff - Gradient of exp", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  const Variable f = suboptimal::exp(x * y);
+  const Vector2v grad = gradient(f, Vector2v{x, y});
+
+  const double x_val = GENERATE(take(10, random(-10.0, 10.0)));
+  const double y_val = GENERATE(take(10, random(-10.0, 10.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+
+  CHECK(suboptimal::getValues(grad).isApprox(Eigen::Vector2d{y_val * std::exp(x_val * y_val),  //
+                                                             x_val * std::exp(x_val * y_val)}));
+}
+
+TEST_CASE("Autodiff - Gradient of log", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  const Variable f = suboptimal::log(x * y);
+  const Vector2v grad = gradient(f, Vector2v{x, y});
+
+  const double x_val = GENERATE(take(10, random(0.1, 100.0)));
+  const double y_val = GENERATE(take(10, random(0.1, 100.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+
+  CHECK(suboptimal::getValues(grad).isApprox(Eigen::Vector2d{1 / x_val,  //
+                                                             1 / y_val}));
+}
+
+TEST_CASE("Autodiff - Gradient of pow", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  const Variable f = suboptimal::pow(x, y);
+  const Vector2v grad = gradient(f, Vector2v{x, y});
+
+  const double x_val = GENERATE(take(10, random(0.0, 100.0)));
+  const double y_val = GENERATE(take(10, random(-100.0, 100.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+
+  CHECK(suboptimal::getValues(grad).isApprox(Eigen::Vector2d{y_val * std::pow(x_val, y_val - 1),  //
+                                                             std::log(x_val) * std::pow(x_val, y_val)}));
+}
+
+TEST_CASE("Autodiff - Gradient of hypot", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  const Variable f = suboptimal::hypot(x, y);
+  const Vector2v grad = gradient(f, Vector2v{x, y});
+
+  const double x_val = GENERATE(take(10, random(-100.0, 100.0)));
+  const double y_val = GENERATE(take(10, random(-100.0, 100.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+
+  CHECK(suboptimal::getValues(grad).isApprox(Eigen::Vector2d{x_val / std::hypot(x_val, y_val),  //
+                                                             y_val / std::hypot(x_val, y_val)}));
+}
+
+TEST_CASE("Autodiff - Gradients of trig functions", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  Variable z{};
+  const Variable f = suboptimal::sin(x) * suboptimal::cos(y) * suboptimal::tan(z);
+  const Vector3v grad = gradient(f, Vector3v{x, y, z});
+
+  const double x_val = GENERATE(take(5, random(-100.0, 100.0)));
+  const double y_val = GENERATE(take(5, random(-100.0, 100.0)));
+  const double z_val = GENERATE(take(5, random(-100.0, 100.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+  z.setValue(z_val);
+
+  CHECK(suboptimal::getValues(grad).isApprox(
+      Eigen::Vector3d{std::cos(y_val) * std::tan(z_val) * std::cos(x_val),   //
+                      -std::sin(x_val) * std::sin(y_val) * std::tan(z_val),  //
+                      std::sin(x_val) * std::cos(y_val) / (std::cos(z_val) * std::cos(z_val))}));
+}
+
+TEST_CASE("Autodiff - Gradients of inverse trig functions", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  Variable z{};
+  const Variable f = suboptimal::asin(x) * suboptimal::acos(y) * suboptimal::atan(z);
+  const Vector3v grad = gradient(f, Vector3v{x, y, z});
+
+  const double x_val = GENERATE(take(5, random(-1.0, 1.0)));
+  const double y_val = GENERATE(take(5, random(-1.0, 1.0)));
+  const double z_val = GENERATE(take(5, random(-100.0, 100.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+  z.setValue(z_val);
+
+  CHECK(suboptimal::getValues(grad).isApprox(
+      Eigen::Vector3d{std::acos(y_val) * std::atan(z_val) / std::sqrt(1 - x_val * x_val),   //
+                      -std::asin(x_val) * std::atan(z_val) / std::sqrt(1 - y_val * y_val),  //
+                      std::asin(x_val) * std::acos(y_val) / (1 + z_val * z_val)}));
+}
+
+TEST_CASE("Autodiff - Gradient of atan2", "[autodiff]") {
+  Variable x{};
+  Variable y{};
+  const Variable f = suboptimal::atan2(y, x);
+  const Vector2v grad = gradient(f, Vector2v{x, y});
+
+  const double x_val = GENERATE(take(10, random(-100.0, 100.0)));
+  const double y_val = GENERATE(take(10, random(-100.0, 100.0)));
+  x.setValue(x_val);
+  y.setValue(y_val);
+
+  CHECK(suboptimal::getValues(grad).isApprox(Eigen::Vector2d{-y_val / (x_val * x_val + y_val * y_val),  //
+                                                             x_val / (x_val * x_val + y_val * y_val)}));
+}