Add autodiff support

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,48 @@
+# 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.
+
+option(SPECTRE_AUTODIFF "Enable automatic differentiation" OFF)
+
+if (NOT SPECTRE_AUTODIFF)
+  return()
+endif()
+
+find_package(autodiff QUIET)
+
+if (NOT autodiff_FOUND)
+  if (NOT SPECTRE_FETCH_MISSING_DEPS)
+    message(FATAL_ERROR "Could not find autodiff. If you want to fetch "
+      "missing dependencies automatically, set SPECTRE_FETCH_MISSING_DEPS=ON.")
+    return()
+  endif()
+  message(STATUS "Fetching autodiff")
+  include(FetchContent)
+  FetchContent_Declare(autodiff
+    GIT_REPOSITORY https://github.com/autodiff/autodiff/
+    GIT_TAG v1.1.2
+    GIT_SHALLOW TRUE
+    ${SPECTRE_FETCHCONTENT_BASE_ARGS}
+  )
+  set(AUTODIFF_BUILD_TESTS OFF CACHE BOOL "Build autodiff tests")
+  set(AUTODIFF_BUILD_EXAMPLES OFF CACHE BOOL "Build autodiff examples")
+  set(AUTODIFF_BUILD_PYTHON OFF CACHE BOOL "Build autodiff Python bindings")
+  set(AUTODIFF_BUILD_DOCS OFF CACHE BOOL "Build autodiff documentation")
+  FetchContent_MakeAvailable(autodiff)
+endif()
+
+add_compile_definitions(SPECTRE_AUTODIFF)
+
+set_property(
+  GLOBAL APPEND PROPERTY SPECTRE_THIRD_PARTY_LIBS
+  autodiff::autodiff
+  )

docs/Installation/BuildSystem.md

@@ -178,6 +178,8 @@ alphabetical order):
   - Set to a path to a SpEC installation (the SpEC repository root) to link in
     SpEC libraries. In particular, the SpEC::Exporter library is linked in and
     enables loading SpEC data into SpECTRE. See \ref installation for details.
+- SPECTRE_AUTODIFF
+  - Enable automatic differentation (default is `OFF`). This is experimental.
 - SPECTRE_DEBUG
   - Defines `SPECTRE_DEBUG` macro to enable `ASSERT`s and other debug
     checks so they can be used in Release builds. That is, you get sanity checks

docs/Installation/Installation.md

@@ -255,6 +255,8 @@ The following dependencies will be fetched automatically if you set
   matplotlib
 * [xsimd](https://github.com/xtensor-stack/xsimd) 11.0.1 or newer - for manual
   vectorization
+* [autodiff](https://github.com/autodiff/autodiff/) 1.1.2 or newer - for
+  automatic differentiation
 * [libbacktrace](https://github.com/ianlancetaylor/libbacktrace) - to show
   source files and line numbers in backtraces of errors and asserts. Available
   by default on many systems, so you may not have to install it at all. The

src/DataStructures/CMakeLists.txt

@@ -83,6 +83,10 @@ target_link_libraries(
   Utilities
   )
 
+if (TARGET autodiff::autodiff)
+  target_link_libraries(${LIBRARY} PUBLIC autodiff::autodiff)
+endif()
+
 add_subdirectory(Blaze)
 add_subdirectory(DataBox)
 add_subdirectory(Python)

src/DataStructures/DynamicVector.hpp

@@ -8,6 +8,10 @@
 
 #pragma once
 
+#ifdef SPECTRE_AUTODIFF
+#include <autodiff/common/numbertraits.hpp>
+#include <autodiff/common/vectortraits.hpp>
+#endif  // SPECTRE_AUTODIFF
 #include <blaze/math/DynamicVector.h>
 #include <pup.h>
 #include <type_traits>
@@ -94,3 +98,25 @@ struct Options::create_from_yaml<blaze::DynamicVector<T, TF, Tag>> {
     return result;
   }
 };
+
+#ifdef SPECTRE_AUTODIFF
+namespace autodiff {
+namespace detail {
+
+template <typename T, bool TF /* , typename Alloc, typename Tag */>
+struct VectorTraits<blaze::DynamicVector<T, TF>> {
+  using ValueType = T;
+
+  template <typename NewValueType>
+  using ReplaceValueType = blaze::DynamicVector<NewValueType, TF>;
+};
+
+template <typename T, bool TF /* , typename Alloc, typename Tag */>
+struct NumberTraits<blaze::DynamicVector<T, TF>> {
+  using NumericType = typename NumberTraits<T>::NumericType;
+  static constexpr auto Order = NumberTraits<T>::Order;
+};
+
+}  // namespace detail
+}  // namespace autodiff
+#endif  // SPECTRE_AUTODIFF

src/DataStructures/Tensor/Expressions/DataTypeSupport.hpp

@@ -59,10 +59,7 @@ constexpr bool is_supported_number_datatype_v =
 ///
 /// \tparam X the `Tensor` data type
 template <typename X>
-struct is_supported_tensor_datatype
-    : std::disjunction<
-          std::is_same<X, double>, std::is_same<X, std::complex<double>>,
-          std::is_same<X, DataVector>, std::is_same<X, ComplexDataVector>> {};
+struct is_supported_tensor_datatype : std::true_type {};
 
 template <typename X>
 constexpr bool is_supported_tensor_datatype_v =

src/DataStructures/Tensor/Tensor.hpp

@@ -103,17 +103,17 @@ class Tensor<X, Symm, IndexList<Indices...>> {
                 "If you are sure you need rank 5 or higher Tensor's please "
                 "file an issue on GitHub or discuss with a core developer of "
                 "SpECTRE.");
-  static_assert(
-      std::is_same_v<X, std::complex<double>> or std::is_same_v<X, double> or
-          std::is_same_v<X, ComplexDataVector> or
-          std::is_same_v<X, ComplexModalVector> or
-          std::is_same_v<X, DataVector> or std::is_same_v<X, ModalVector> or
-          is_spin_weighted_of_v<ComplexDataVector, X> or
-          is_spin_weighted_of_v<ComplexModalVector, X> or
-          simd::is_batch<X>::value,
-      "Unsupported type. While other types are technically possible it is not "
-      "clear that Tensor is the correct container for them. Please seek advice "
-      "on the topic by discussing with the SpECTRE developers.");
+  // static_assert(
+  //     std::is_same_v<X, std::complex<double>> or std::is_same_v<X, double> or
+  //         std::is_same_v<X, ComplexDataVector> or
+  //         std::is_same_v<X, ComplexModalVector> or
+  //         std::is_same_v<X, DataVector> or std::is_same_v<X, ModalVector> or
+  //         is_spin_weighted_of_v<ComplexDataVector, X> or
+  //         is_spin_weighted_of_v<ComplexModalVector, X> or
+  //         simd::is_batch<X>::value,
+  //     "Unsupported type. While other types are technically possible it is not
+  //     " "clear that Tensor is the correct container for them. Please seek
+  //     advice " "on the topic by discussing with the SpECTRE developers.");
 
  public:
   /// The type of the sequence that holds the data

src/Domain/CoordinateMaps/CMakeLists.txt

@@ -97,5 +97,9 @@ target_link_libraries(
   SphericalHarmonics
   )
 
+if (TARGET autodiff::autodiff)
+  target_link_libraries(${LIBRARY} PRIVATE autodiff::autodiff)
+endif()
+
 add_subdirectory(Python)
 add_subdirectory(TimeDependent)

src/Domain/CoordinateMaps/Interval.cpp

@@ -253,18 +253,34 @@ 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>))
+GENERATE_INSTANTIATIONS(INSTANTIATE_JAC,
+                        (double, DataVector,
+                         std::reference_wrapper<const double>,
+                         std::reference_wrapper<const DataVector>))
+
+#ifdef SPECTRE_AUTODIFF
+#include <autodiff/forward/dual.hpp>
+#include <autodiff/forward/real.hpp>
+#include <autodiff/reverse/var.hpp>
+GENERATE_INSTANTIATIONS(INSTANTIATE,
+                        (autodiff::dual, autodiff::real, autodiff::var))
+#endif  // SPECTRE_AUTODIFF
+
 #undef DTYPE
 #undef INSTANTIATE
+#undef INSTANTIATE_JAC
 }  // namespace domain::CoordinateMaps

src/Domain/CoordinateMaps/Wedge.cpp

@@ -251,17 +251,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 +510,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 +806,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,24 +1022,39 @@ 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>))
+GENERATE_INSTANTIATIONS(INSTANTIATE_DTYPE_JAC, (2, 3),
+                        (double, DataVector,
+                         std::reference_wrapper<const double>,
+                         std::reference_wrapper<const DataVector>))
+
+#ifdef SPECTRE_AUTODIFF
+#include <autodiff/forward/dual.hpp>
+#include <autodiff/forward/real.hpp>
+#include <autodiff/reverse/var.hpp>
+GENERATE_INSTANTIATIONS(INSTANTIATE_DTYPE, (2, 3),
+                        (autodiff::dual, autodiff::real, autodiff::var))
+#endif  // SPECTRE_AUTODIFF
 
 #undef DIM
 #undef DTYPE
 #undef INSTANTIATE_DIM
 #undef INSTANTIATE_DTYPE
+#undef INSTANTIATE_DTYPE_JAC
 }  // namespace domain::CoordinateMaps

src/Domain/Structure/CMakeLists.txt

@@ -66,3 +66,7 @@ target_link_libraries(
   Spectral
   Utilities
   )
+
+if (TARGET autodiff::autodiff)
+  target_link_libraries(${LIBRARY} PRIVATE autodiff::autodiff)
+endif()

src/Domain/Structure/OrientationMap.cpp

@@ -292,6 +292,14 @@ GENERATE_INSTANTIATIONS(INSTANTIATION, (1, 2, 3),
                          std::reference_wrapper<const double> const,
                          std::reference_wrapper<const DataVector> const))
 
+#ifdef SPECTRE_AUTODIFF
+#include <autodiff/forward/dual.hpp>
+#include <autodiff/forward/real.hpp>
+#include <autodiff/reverse/var.hpp>
+GENERATE_INSTANTIATIONS(INSTANTIATION, (1, 2, 3),
+                        (autodiff::dual, autodiff::real, autodiff::var))
+#endif  // SPECTRE_AUTODIFF
+
 #undef INSTANTIATION
 #undef DTYPE
 #undef DIM

tests/Unit/DataStructures/CMakeLists.txt

@@ -47,6 +47,10 @@ set(LIBRARY_SOURCES
   Test_VectorImplTestHelper.cpp
   )
 
+if (TARGET autodiff::autodiff)
+  list(APPEND LIBRARY_SOURCES Test_Autodiff.cpp)
+endif()
+
 add_subdirectory(DataBox)
 add_subdirectory(Tensor)
 
@@ -62,6 +66,10 @@ target_link_libraries(
   Utilities
   )
 
+if (TARGET autodiff::autodiff)
+  target_link_libraries(${LIBRARY} PRIVATE autodiff::autodiff)
+endif()
+
 # [example_add_pybindings_test]
 spectre_add_python_bindings_test(
   "Unit.DataStructures.Python.DataVector"