WIP autodiff

CMakeLists.txt

@@ -154,6 +154,7 @@ include(SetupLapack)
 include(SetupLIBXSMM)
 include(SetupGsl)
 
+include(SetupAutodiff)
 include(SetupBlaze)
 include(SetupFuka)
 include(SetupGoogleBenchmark)

cmake/SetupAutodiff.cmake

@@ -0,0 +1,14 @@
+# Distributed under the MIT License.
+# See LICENSE.txt for details.
+
+# Find autodiff (https://github.com/autodiff/autodiff)
+#
+# Note that Boost also has an autodiff module, but it uses templates to
+# distinguish multiple variables. This means a std::array or Tensor can't hold
+# the Boost autodiff variables. For example, to evaluate a 2D function f(x, y),
+# the type for x would be fvar<double> and the type for y would be the nested
+# fvar<fvar<double>>. On the other hand, the autodiff library works with a
+# simpler autodiff::dual or autodiff::var type, which can be stored in a
+# std::array or Tensor.
+
+find_package(autodiff REQUIRED)

src/Domain/CoordinateMaps/CMakeLists.txt

@@ -83,6 +83,7 @@ spectre_target_headers(
 target_link_libraries(
   ${LIBRARY}
   PRIVATE
+  autodiff::autodiff
   RootFinding
   PUBLIC
   Boost::boost

src/Domain/CoordinateMaps/Interval.cpp

@@ -3,6 +3,8 @@
 
 #include "Domain/CoordinateMaps/Interval.hpp"
 
+#include <autodiff/forward/dual.hpp>
+#include <autodiff/reverse/var.hpp>
 #include <cmath>
 #include <optional>
 #include <pup.h>
@@ -253,18 +255,25 @@ bool operator==(const CoordinateMaps::Interval& lhs,
 // Explicit instantiations
 #define DTYPE(data) BOOST_PP_TUPLE_ELEM(0, data)
 
-#define INSTANTIATE(_, data)                                                   \
-  template std::array<tt::remove_cvref_wrap_t<DTYPE(data)>, 1>                 \
-  Interval::operator()(const std::array<DTYPE(data), 1>& source_coords) const; \
-  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, 1, Frame::NoFrame>   \
-  Interval::jacobian(const std::array<DTYPE(data), 1>& source_coords) const;   \
-  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, 1, Frame::NoFrame>   \
-  Interval::inv_jacobian(const std::array<DTYPE(data), 1>& source_coords)      \
+#define INSTANTIATE(_, data)                                   \
+  template std::array<tt::remove_cvref_wrap_t<DTYPE(data)>, 1> \
+  Interval::operator()(const std::array<DTYPE(data), 1>& source_coords) const;
+
+#define INSTANTIATE_JAC(_, data)                                             \
+  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, 1, Frame::NoFrame> \
+  Interval::jacobian(const std::array<DTYPE(data), 1>& source_coords) const; \
+  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, 1, Frame::NoFrame> \
+  Interval::inv_jacobian(const std::array<DTYPE(data), 1>& source_coords)    \
       const;
 
 GENERATE_INSTANTIATIONS(INSTANTIATE, (double, DataVector,
                                       std::reference_wrapper<const double>,
-                                      std::reference_wrapper<const DataVector>))
+                                      std::reference_wrapper<const DataVector>,
+                                      autodiff::dual, autodiff::var))
+GENERATE_INSTANTIATIONS(INSTANTIATE_JAC,
+                        (double, DataVector,
+                         std::reference_wrapper<const double>,
+                         std::reference_wrapper<const DataVector>))
 #undef DTYPE
 #undef INSTANTIATE
 }  // namespace domain::CoordinateMaps

src/Domain/CoordinateMaps/Wedge.cpp

@@ -3,6 +3,8 @@
 
 #include "Domain/CoordinateMaps/Wedge.hpp"
 
+#include <autodiff/forward/dual.hpp>
+#include <autodiff/reverse/var.hpp>
 #include <climits>
 #include <cmath>
 #include <cstddef>
@@ -251,17 +253,17 @@ template <bool FuncIsXi, typename T>
 tt::remove_cvref_wrap_t<T> Wedge<Dim>::get_cap_angular_function(
     const T& lowercase_xi_or_eta) const {
   constexpr auto cap_index = static_cast<size_t>(not FuncIsXi);
+  if (not with_equiangular_map_) {
+    return lowercase_xi_or_eta;
+  }
   if (opening_angles_.has_value() and
       opening_angles_distribution_.has_value()) {
-    return with_equiangular_map_
-               ? tan(0.5 * opening_angles_.value()[cap_index]) *
-                     tan(0.5 * opening_angles_distribution_.value()[cap_index] *
-                         lowercase_xi_or_eta) /
-                     tan(0.5 * opening_angles_distribution_.value()[cap_index])
-               : lowercase_xi_or_eta;
+    return tan(0.5 * opening_angles_.value()[cap_index]) *
+           tan(0.5 * opening_angles_distribution_.value()[cap_index] *
+               lowercase_xi_or_eta) /
+           tan(0.5 * opening_angles_distribution_.value()[cap_index]);
   } else {
-    return with_equiangular_map_ ? tan(M_PI_4 * lowercase_xi_or_eta)
-                                 : lowercase_xi_or_eta;
+    return tan(M_PI_4 * lowercase_xi_or_eta);
   }
 }
 
@@ -510,11 +512,14 @@ std::array<tt::remove_cvref_wrap_t<T>, Dim> Wedge<Dim>::operator()(
   const bool zero_offset = not cube_half_length_.has_value();
 
   std::array<ReturnType, Dim> physical_coords{};
-  physical_coords[radial_coord] =
-      zero_offset ? generalized_z
-                  : generalized_z * (1.0 - rotated_focus[radial_coord] /
-                                               cube_half_length_.value()) +
-                        rotated_focus[radial_coord];
+  if (zero_offset) {
+    physical_coords[radial_coord] = generalized_z;
+  } else {
+    physical_coords[radial_coord] =
+        generalized_z *
+            (1.0 - rotated_focus[radial_coord] / cube_half_length_.value()) +
+        rotated_focus[radial_coord];
+  }
   if (zero_offset) {
     physical_coords[polar_coord] = generalized_z * cap[0];
   } else {
@@ -803,7 +808,6 @@ Wedge<Dim>::inv_jacobian(const std::array<T, Dim>& source_coords) const {
     xi *= 0.5;
   }
 
-
   std::array<ReturnType, Dim - 1> cap{};
   std::array<ReturnType, Dim - 1> cap_deriv{};
   cap[0] = get_cap_angular_function<true>(xi);
@@ -1020,18 +1024,25 @@ bool operator!=(const Wedge<Dim>& lhs, const Wedge<Dim>& rhs) {
 #define INSTANTIATE_DTYPE(_, data)                                     \
   template std::array<tt::remove_cvref_wrap_t<DTYPE(data)>, DIM(data)> \
   Wedge<DIM(data)>::operator()(                                        \
-      const std::array<DTYPE(data), DIM(data)>& source_coords) const;  \
-  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, DIM(data),   \
-                    Frame::NoFrame>                                    \
-  Wedge<DIM(data)>::jacobian(                                          \
-      const std::array<DTYPE(data), DIM(data)>& source_coords) const;  \
-  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, DIM(data),   \
-                    Frame::NoFrame>                                    \
-  Wedge<DIM(data)>::inv_jacobian(                                      \
+      const std::array<DTYPE(data), DIM(data)>& source_coords) const;
+
+#define INSTANTIATE_DTYPE_JAC(_, data)                                \
+  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, DIM(data),  \
+                    Frame::NoFrame>                                   \
+  Wedge<DIM(data)>::jacobian(                                         \
+      const std::array<DTYPE(data), DIM(data)>& source_coords) const; \
+  template tnsr::Ij<tt::remove_cvref_wrap_t<DTYPE(data)>, DIM(data),  \
+                    Frame::NoFrame>                                   \
+  Wedge<DIM(data)>::inv_jacobian(                                     \
       const std::array<DTYPE(data), DIM(data)>& source_coords) const;
 
 GENERATE_INSTANTIATIONS(INSTANTIATE_DIM, (2, 3))
 GENERATE_INSTANTIATIONS(INSTANTIATE_DTYPE, (2, 3),
+                        (double, DataVector,
+                         std::reference_wrapper<const double>,
+                         std::reference_wrapper<const DataVector>,
+                         autodiff::dual, autodiff::var))
+GENERATE_INSTANTIATIONS(INSTANTIATE_DTYPE_JAC, (2, 3),
                         (double, DataVector,
                          std::reference_wrapper<const double>,
                          std::reference_wrapper<const DataVector>))

src/Domain/Structure/OrientationMap.cpp

@@ -3,6 +3,8 @@
 
 #include "Domain/Structure/OrientationMap.hpp"
 
+#include <autodiff/forward/dual.hpp>
+#include <autodiff/reverse/var.hpp>
 #include <functional>
 #include <ostream>
 #include <pup.h>
@@ -261,12 +263,12 @@ tnsr::Ij<double, VolumeDim, Frame::NoFrame> discrete_rotation_inverse_jacobian(
 
 #define DIM(data) BOOST_PP_TUPLE_ELEM(0, data)
 
-#define INSTANTIATION(r, data)                                              \
-  template class OrientationMap<DIM(data)>;                                 \
-  template tnsr::Ij<double, DIM(data), Frame::NoFrame>                      \
-  discrete_rotation_jacobian(const OrientationMap<DIM(data)>& orientation); \
-  template tnsr::Ij<double, DIM(data), Frame::NoFrame>                      \
-  discrete_rotation_inverse_jacobian(                                       \
+#define INSTANTIATION(r, data)                                                \
+  template class OrientationMap<DIM(data)>;                                   \
+  template tnsr::Ij<double, DIM(data), Frame::NoFrame>                        \
+  discrete_rotation_jacobian(const OrientationMap<DIM(data)>& orientation);   \
+  template tnsr::Ij<double, DIM(data), Frame::NoFrame>                        \
+  discrete_rotation_inverse_jacobian(                                         \
       const OrientationMap<DIM(data)>& orientation);
 
 GENERATE_INSTANTIATIONS(INSTANTIATION, (1, 2, 3))
@@ -290,7 +292,8 @@ GENERATE_INSTANTIATIONS(INSTANTIATION, (1, 2, 3),
                          std::reference_wrapper<double> const,
                          std::reference_wrapper<DataVector> const,
                          std::reference_wrapper<const double> const,
-                         std::reference_wrapper<const DataVector> const))
+                         std::reference_wrapper<const DataVector> const,
+                         autodiff::dual, autodiff::var))
 
 #undef INSTANTIATION
 #undef DTYPE

src/Utilities/CMakeLists.txt

@@ -90,6 +90,7 @@ spectre_target_headers(
 target_link_libraries(
   ${LIBRARY}
   PUBLIC
+  autodiff::autodiff
   BLAS::BLAS
   Blaze
   Boost::boost

tests/Unit/Domain/CoordinateMaps/CMakeLists.txt

@@ -4,32 +4,8 @@
 set(LIBRARY "Test_CoordinateMaps")
 
 set(LIBRARY_SOURCES
-  Test_Affine.cpp
-  Test_BulgedCube.cpp
-  Test_Composition.cpp
-  Test_CoordinateMap.cpp
-  Test_CylindricalEndcap.cpp
-  Test_CylindricalFlatEndcap.cpp
-  Test_CylindricalFlatSide.cpp
-  Test_CylindricalSide.cpp
-  Test_DiscreteRotation.cpp
-  Test_Distribution.cpp
-  Test_EquatorialCompression.cpp
-  Test_Equiangular.cpp
-  Test_Frustum.cpp
-  Test_Identity.cpp
   Test_Interval.cpp
-  Test_KerrHorizonConforming.cpp
-  Test_ProductMaps.cpp
-  Test_Rotation.cpp
-  Test_SpecialMobius.cpp
-  Test_Tags.cpp
-  Test_TimeDependentHelpers.cpp
-  Test_UniformCylindricalEndcap.cpp
-  Test_UniformCylindricalFlatEndcap.cpp
-  Test_UniformCylindricalSide.cpp
   Test_Wedge2D.cpp
-  Test_Wedge3D.cpp
   )
 
 add_test_library(${LIBRARY} "${LIBRARY_SOURCES}")
@@ -49,5 +25,3 @@ target_link_libraries(
   Spectral
   Utilities
 )
-
-add_subdirectory(TimeDependent)

tests/Unit/Domain/CoordinateMaps/Test_Interval.cpp

@@ -30,6 +30,7 @@ void test_coordinate_map(
     test_inverse_map(map, random_point);
     test_jacobian(map, random_point);
     test_inv_jacobian(map, random_point);
+    test_jacobian_autodiff(map, random_point);
   }
   const std::array<double, 1> point_xA{{xA}};
   const std::array<double, 1> point_xB{{xB}};

tests/Unit/Elliptic/Systems/Punctures/CMakeLists.txt

@@ -13,6 +13,7 @@ add_test_library(${LIBRARY} "${LIBRARY_SOURCES}")
 target_link_libraries(
   ${LIBRARY}
   PRIVATE
+  autodiff::autodiff
   DataStructuresHelpers
   Punctures
   )

tests/Unit/Elliptic/Systems/Punctures/Test_Equations.cpp

@@ -4,6 +4,7 @@
 #include "Framework/TestingFramework.hpp"
 
 #include <array>
+#include <autodiff/reverse/var.hpp>
 #include <cstddef>
 #include <string>
 
@@ -18,6 +19,35 @@
 
 namespace {
 
+template <typename DataType, typename VarDataType>
+void add_sources(const gsl::not_null<VarDataType*> puncture_equation,
+                 const DataType& alpha, const DataType& beta,
+                 const VarDataType& field) {
+  *puncture_equation -= beta / pow<7>(alpha * (field + 1.) + 1.);
+}
+
+void add_linearized_sources_autodiff(
+    const gsl::not_null<Scalar<DataVector>*> linearized_puncture_equation,
+    const Scalar<DataVector>& alpha, const Scalar<DataVector>& beta,
+    const Scalar<DataVector>& field,
+    const Scalar<DataVector>& field_correction) {
+  // Works:
+  // - Going over each element of the DataVector, evaluating autodiff pointwise
+  // TODO:
+  // - Vectorize the autodiff evaluation, e.g. by using a new type
+  //   DualDataVector = VectorImpl<autodiff::dual> (or the reverse mode
+  //   equivalent)
+  // - Add support for Tensor<DualDataVector> (probably just extend the
+  //   static_assert to allow this).
+  for (size_t i = 0; i < get(field).size(); ++i) {
+    const autodiff::var field_i{get(field)[i]};
+    autodiff::var result{get(*linearized_puncture_equation)[i]};
+    add_sources(make_not_null(&result), get(alpha)[i], get(beta)[i], field_i);
+    const auto dS_du = autodiff::derivatives(result, autodiff::wrt(field_i))[0];
+    get(*linearized_puncture_equation)[i] += dS_du * get(field_correction)[i];
+  }
+}
+
 void test_equations(const DataVector& used_for_size) {
   const double eps = 1.e-12;
   const auto seed = std::random_device{}();
@@ -30,6 +60,10 @@ void test_equations(const DataVector& used_for_size) {
       &Punctures::add_linearized_sources, "Equations", {"linearized_sources"},
       {{{-1., 1.}, {-1., 1.}, {-1., 1.}, {-1., 1.}}}, used_for_size, eps, seed,
       fill_result_tensors);
+  pypp::check_with_random_values<4>(
+      &add_linearized_sources_autodiff, "Equations", {"linearized_sources"},
+      {{{-1., 1.}, {-1., 1.}, {-1., 1.}, {-1., 1.}}}, used_for_size, eps, seed,
+      fill_result_tensors);
 }
 
 void test_computers(const DataVector& used_for_size) {

tests/Unit/Helpers/Domain/CMakeLists.txt

@@ -12,6 +12,7 @@ add_spectre_library(${LIBRARY} ${SPECTRE_TEST_LIBS_TYPE} ${LIBRARY_SOURCES})
 target_link_libraries(
   ${LIBRARY}
   INTERFACE
+  autodiff::autodiff
   Boost::boost
   CoordinateMaps
   FunctionsOfTime

tests/Unit/Helpers/Domain/CoordinateMaps/TestMapHelpers.hpp

@@ -10,6 +10,7 @@
 
 #include <algorithm>
 #include <array>
+#include <autodiff/reverse/var.hpp>
 #include <cmath>
 #include <functional>
 #include <limits>
@@ -120,6 +121,36 @@ void test_jacobian(const Map& map,
   }
 }
 
+template <typename Map>
+void test_jacobian_autodiff(const Map& map,
+                            const std::array<double, Map::dim>& test_point) {
+  INFO("Test Jacobian with autodiff");
+  CAPTURE(test_point);
+  constexpr size_t Dim = Map::dim;
+  // Use reverse mode because we want the full Jacobian (all derivatives)
+  std::array<autodiff::var, Dim> x{};
+  for (size_t i = 0; i < Dim; ++i) {
+    gsl::at(x, i) = gsl::at(test_point, i);
+  }
+  const auto y = map(x);
+  const auto jacobian = map.jacobian(test_point);
+  for (size_t i = 0; i < Dim; ++i) {
+    std::array<double, Dim> deriv_i{};
+    if constexpr (Dim == 1) {
+      deriv_i = autodiff::derivatives(gsl::at(y, i), autodiff::wrt(x[0]));
+    } else if constexpr (Dim == 2) {
+      deriv_i = autodiff::derivatives(gsl::at(y, i), autodiff::wrt(x[0], x[1]));
+    } else if constexpr (Dim == 3) {
+      deriv_i =
+          autodiff::derivatives(gsl::at(y, i), autodiff::wrt(x[0], x[1], x[2]));
+    }
+    for (size_t j = 0; j < Dim; ++j) {
+      INFO("i: " << i << " j: " << j);
+      CHECK(jacobian.get(i, j) == approx(gsl::at(deriv_i, j)));
+    }
+  }
+}
+
 template <typename Map>
 void test_jacobian(
     const Map& map, const std::array<double, Map::dim>& test_point,
@@ -586,16 +617,19 @@ void test_suite_for_map_on_unit_cube(const Map& map) {
     test_coordinate_map_argument_types(map_to_test, origin);
 
     test_jacobian(map_to_test, origin);
+    test_jacobian_autodiff(map_to_test, origin);
     test_inv_jacobian(map_to_test, origin);
     test_inverse_map(map_to_test, origin);
 
     for (VolumeCornerIterator<Map::dim> vci{}; vci; ++vci) {
       test_jacobian(map_to_test, vci.coords_of_corner());
+      test_jacobian_autodiff(map_to_test, vci.coords_of_corner());
       test_inv_jacobian(map_to_test, vci.coords_of_corner());
       test_inverse_map(map_to_test, vci.coords_of_corner());
     }
 
     test_jacobian(map_to_test, random_point);
+    test_jacobian_autodiff(map_to_test, random_point);
     test_inv_jacobian(map_to_test, random_point);
     test_inverse_map(map_to_test, random_point);
   };