Place guards correctly on my additions so that miniapp will still compile without Enzyme

The directional derivative was already fixed in an earlier commit.

linalg/dual.hpp

@@ -376,11 +376,11 @@ std::ostream& operator<<(std::ostream& os, dual<value_type, gradient_type> A)
 MFEM_HOST_DEVICE constexpr dual<double, double> make_dual(double x) { return {x, 1.0}; }
 
 /** @brief return the "value" part from a given type. For non-dual types, this is just the identity function */
-template <typename T> MFEM_HOST_DEVICE T get_value(const T& arg) { return arg; }
+MFEM_HOST_DEVICE constexpr double get_value(const double& arg) { return arg; }
 
 /** @brief return the "value" part from a dual number type */
 template <typename value_type, typename gradient_type>
-MFEM_HOST_DEVICE gradient_type get_value(dual<value_type, gradient_type> arg)
+MFEM_HOST_DEVICE value_type get_value(dual<value_type, gradient_type> arg)
 {
    return arg.value;
 }

linalg/tensor.hpp

@@ -1582,7 +1582,7 @@ tensor<T, n, n> inv(const tensor<T, n, n>& A)
  * TODO: compare performance of this hardcoded implementation to just using inv() directly
  */
 template <typename value_type, typename gradient_type, int n> MFEM_HOST_DEVICE
-dual<value_type, gradient_type> inv(
+tensor<dual<value_type, gradient_type>, n, n> inv(
    tensor<dual<value_type, gradient_type>, n, n> A)
 {
    auto invA = inv(get_value(A));
@@ -1601,6 +1601,29 @@ dual<value_type, gradient_type> inv(
    });
 }
 
+/** @brief  */
+template <typename value_type, typename gradient_type, int m>
+MFEM_HOST_DEVICE tensor< value_type, m> get_value(const tensor<dual<value_type, gradient_type>, m>& arg)
+{
+  tensor<value_type, m> value{};
+  for (int i = 0; i < m; i++) {
+    value[i] = arg[i].value;
+  }
+  return value;
+}
+
+template <typename value_type, typename gradient_type, int m, int n>
+MFEM_HOST_DEVICE tensor< value_type, m, n > get_value(const tensor<dual<value_type, gradient_type>, m, n>& arg)
+{
+  tensor<value_type, m, n> value{};
+  for (int i = 0; i < m; i++) {
+    for (int j = 0; j < n; j++) {
+      value[i][j] = arg[i][j].value;
+    }
+  }
+  return value;
+}
+
 /**
  * @brief recursively serialize the entries in a tensor to an output stream.
  * Output format uses braces and comma separators to mimic C syntax for multidimensional array

miniapps/elasticity/elasticity_driver.cpp

@@ -104,7 +104,11 @@ int main(int argc, char *argv[])
    // * EnzymeRev
    // * FiniteDiff
    // * DualNumbers
-   const NeoHookeanMaterial<dimension, GradientType::DualNumbers> material{};
+   // The second template argument sets mode for computing the derivative of the stress wrt
+   // the displacement gradient (ie the elasticities tensor).
+   // The third template argument is for computing the derivative of the stress wrt the
+   // displacement gradient (ie the stress tensor).
+   const NeoHookeanMaterial<dimension, GradientType::EnzymeFwd, GradientType::EnzymeRev> material{};
    elasticity_op.SetMaterial(material);
 
    // Define all essential boundaries. In this specific example, this includes

miniapps/elasticity/makefile

@@ -27,7 +27,7 @@ ELASTICITY_SRC = elasticity_driver.cpp\
 	preconditioners/diagonal_preconditioner.cpp
 ELASTICITY_OBJ = $(ELASTICITY_SRC:.cpp=.o)
 
-MINIAPPS = elasticity
+MINIAPPS = elasticity tests
 MFEM_TESTS = elasticity
 
 .SUFFIXES:
@@ -49,6 +49,9 @@ ELASTICITY_INCLUDES = -I.
 elasticity: $(ELASTICITY_OBJ)
 	$(MFEM_CXX) $(MFEM_LINK_FLAGS) -o $@ $(ELASTICITY_OBJ) $(MFEM_LIBS)
 
+tests: tests.o
+	$(MFEM_CXX) $(MFEM_LINK_FLAGS) -o $@ $< $(MFEM_LIBS)
+
 %.o: $(SRC)%.cpp $(MFEM_LIB_FILE) $(CONFIG_MK)
 	$(MFEM_CXX) $(MFEM_FLAGS) $(ELASTICITY_INCLUDES) -c $< -o $@
 

miniapps/elasticity/materials/neohookean.hpp

@@ -25,7 +25,10 @@ using mfem::internal::make_tensor;
  *
  * Defines a Neo-Hookean material response. It satisfies the material_type
  * interface for ElasticityOperator::SetMaterial. This material type allows
- * choosing the method of derivative calculation in `action_of_gradient`.
+ * choosing the method of derivative calculation for the stress (which is
+ * the derivative of the strain energy density) and for the 
+ * `action_of_gradient`, which applies the tangent elastic tensor operator
+ * as a directional derivative.
  * Choices include methods derived by hand using symbolic calculation and a
  * variety of automatically computed gradient applications, like
  * - Enzyme forward mode
@@ -34,13 +37,35 @@ using mfem::internal::make_tensor;
  * - Finite difference mode
  *
  * @tparam dim
- * @tparam gradient_type
+ * @tparam stress_gradient_type How to compute the derivative of the stress
+ * @tparam energy_gradient_type How to compute the derivative of the energy, which is the stress
  */
-template <int dim = 3, GradientType gradient_type = GradientType::Symbolic>
+template <int dim = 3, GradientType stress_gradient_type = GradientType::Symbolic, GradientType energy_gradient_type = GradientType::Symbolic>
 struct NeoHookeanMaterial
 {
    static_assert(dim == 3, "NeoHookean model only defined in 3D");
 
+   /**
+    * @brief Compute strain energy density
+    * @param[in] dudx displacement gradient
+    * @return W strain energy density in reference configuration
+    */
+   template <typename T>
+   MFEM_HOST_DEVICE T
+   strain_energy_density(const tensor<T, dim, dim> & dudx) const
+   {
+      // Incorrect answer with IsotropicIdentity and enzyme fwddiff mode.
+      // type deduction is not what I expect - why?
+      // Ask Julian
+      auto I = mfem::internal::Identity<dim>();
+      auto F = I + dudx;
+      auto J = det(F);
+      T Jm23 = pow(J, -2.0/3.0);
+      T Wvol = D1*(J - 1.0)*(J - 1.0);
+      T Wdev = C1*(Jm23*inner(F,F) - 3.0);
+      return Wdev + Wvol;
+   }
+
    /**
     * @brief Compute the stress response.
     *
@@ -49,14 +74,161 @@ struct NeoHookeanMaterial
     */
    template <typename T>
    MFEM_HOST_DEVICE tensor<T, dim, dim>
-   stress(const tensor<T, dim, dim> &__restrict__ dudx) const
+   stress(const tensor<T, dim, dim> & dudx) const
    {
-      constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      if (energy_gradient_type == GradientType::Symbolic)
+      {
+         return stress_symbolic(dudx);
+      }
+      else if (energy_gradient_type == GradientType::EnzymeFwd)
+      {
+         return stress_enzyme_fwd(dudx);
+      }
+      else if (energy_gradient_type == GradientType::EnzymeRev)
+      {
+         return stress_enzyme_rev(dudx);
+      }
+      else if (energy_gradient_type == GradientType::FiniteDiff)
+      {
+         return stress_fd(dudx);
+      }
+      else if (energy_gradient_type == GradientType::DualNumbers)
+      {
+         return stress_dual(dudx);
+      }
+   }
+
+   /**
+    * @brief A method to wrap the strain energy density calculation into a static function.
+    *
+    * This is necessary for Enzyme to access the class pointer (self).
+    *
+    * @param[in] self the class pointer
+    * @param[in] dudx the displacement gradient
+    * @param[out] W the strain energy density 
+    */
+   MFEM_HOST_DEVICE static void
+   strain_energy_density_wrapper(NeoHookeanMaterial<dim, stress_gradient_type> *self,
+                                 tensor<double, dim, dim> &dudx, double &W)
+   {
+      W = self->strain_energy_density(dudx);
+   }
+
+   /**
+    * @brief Evaluate the stress symbolically
+    *
+    * @param[in] dudx
+    * @return Piola stress
+    */
+   template <typename T>
+   MFEM_HOST_DEVICE tensor<T, dim, dim>
+   stress_symbolic(const tensor<T, dim, dim> &__restrict__ dudx) const
+   {
+      auto I = mfem::internal::Identity<dim>();
       T J = det(I + dudx);
       T p = -2.0 * D1 * J * (J - 1);
       auto devB = dev(dudx + transpose(dudx) + dot(dudx, transpose(dudx)));
       auto sigma = -(p / J) * I + 2 * (C1 / pow(J, 5.0 / 3.0)) * devB;
-      return sigma;
+      auto F = dudx + I;
+      return J*sigma*transpose(inv(F));
+   }
+
+#ifdef MFEM_USE_ENZYME
+   /**
+    * @brief Evaluate the stress with reverse mode differentiation
+    *
+    * @param[in] dudx
+    * @return Piola stress
+    */
+   template <typename T>
+   MFEM_HOST_DEVICE tensor<T, dim, dim>
+   stress_enzyme_rev(const tensor<T, dim, dim> &dudx) const
+   {
+      T W;
+      T seed{1.0};
+      tensor<T, 3, 3> P{};
+      __enzyme_autodiff<void>(strain_energy_density_wrapper, enzyme_const, this,
+                              enzyme_dup, &dudx, &P, enzyme_dupnoneed, &W, &seed);
+      return P;
+   }
+
+   /**
+    * @brief Evaluate the stress with forward mode differentiation
+    *
+    * @param[in] dudx
+    * @return Piola stress
+    */
+   template <typename T>
+   MFEM_HOST_DEVICE tensor<T, dim, dim>
+   stress_enzyme_fwd(const tensor<T, dim, dim> &dudx) const
+   {
+      tensor<T, dim, dim> direction{};
+      tensor<T, dim, dim> P{};
+      T W;
+      for (int i = 0; i < dim; ++i) {
+         for (int j = 0; j < dim; ++j) {
+            direction[i][j] = 1;
+            __enzyme_fwddiff<void>(strain_energy_density_wrapper, enzyme_const, this,
+                                   enzyme_dupnoneed, &dudx, &direction,
+                                   enzyme_dup, &W, &(P[i][j]));
+            direction[i][j] = 0;
+         }
+      }
+      return P;
+   }
+#endif
+
+   /**
+    * @brief Evaluate the stress with finite differences
+    *
+    * @param[in] dudx
+    * @return Piola stress
+    */
+   template <typename T>
+   MFEM_HOST_DEVICE tensor<T, dim, dim>
+   stress_fd(const tensor<T, dim, dim> &dudx) const
+   {
+      auto H = dudx;
+      auto P = 0.0*dudx;
+      T h{1e-6};
+      for (int i = 0; i < dim; ++i) {
+         for (int j = 0; j < dim; ++j) {
+            H[i][j] += h;
+            auto Wp = strain_energy_density(H);
+            H[i][j] -= 2.0*h;
+            auto Wm = strain_energy_density(H);
+            H[i][j] += h;
+            P[i][j] = (Wp - Wm) / (2.0 * h);
+         }
+      }
+      return P;
+   }
+
+   /**
+    * @brief Evaluate the stress with built-in dual number forward mode
+    *
+    * @param[in] dudx
+    * @return Piola stress
+    */
+   template <typename T>
+   MFEM_HOST_DEVICE tensor<T, dim, dim>
+   stress_dual(const tensor<T, dim, dim> &dudx) const
+   {
+      tensor<T, dim, dim> dir{};
+      tensor<T, dim, dim> P{};
+      for (int k = 0; k < dim; ++k) {
+         for (int l = 0; l < dim; ++l) {
+            dir[k][l] = 1;
+            auto H(make_tensor<dim, dim>([&](int i, int j)
+            {
+               return mfem::internal::dual<T, T> {dudx[i][j], dir[i][j]};
+            }));
+            auto W = strain_energy_density(H);
+            P[k][l] = get_gradient(W);
+            dir[k][l] = 0;
+         }
+      }
+      return P;
    }
 
    /**
@@ -66,11 +238,11 @@ struct NeoHookeanMaterial
     *
     * @param[in] self
     * @param[in] dudx
-    * @param[in] sigma
+    * @param[out] sigma
     * @return stress
     */
    MFEM_HOST_DEVICE static void
-   stress_wrapper(NeoHookeanMaterial<dim, gradient_type> *self,
+   stress_wrapper(NeoHookeanMaterial<dim, stress_gradient_type, energy_gradient_type> *self,
                   tensor<double, dim, dim> &dudx,
                   tensor<double, dim, dim> &sigma)
    {
@@ -93,20 +265,29 @@ struct NeoHookeanMaterial
       constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
 
       tensor<double, dim, dim> F = I + dudx;
-      tensor<double, dim, dim> invF = inv(F);
-      tensor<double, dim, dim> devB =
-         dev(dudx + transpose(dudx) + dot(dudx, transpose(dudx)));
       double J = det(F);
-      double coef = (C1 / pow(J, 5.0 / 3.0));
-      return make_tensor<dim, dim, dim, dim>([&](int i, int j, int k,
-                                                 int l)
-      {
-         return 2.0 * (D1 * J * (i == j) - (5.0 / 3.0) * coef * devB[i][j]) *
-                invF[l][k] +
-                2.0 * coef *
-                ((i == k) * F[j][l] + F[i][l] * (j == k) -
-                 (2.0 / 3.0) * ((i == j) * F[k][l]));
-      });
+      double Jm23 = pow(J, -2.0/3.0);
+      double I1 = inner(F, F);
+      auto invF = inv(F);
+
+      double dWvol = 2.0*D1*(J - 1.0);
+      double ddWvol = 2.0*D1;
+
+      tensor<double, dim, dim, dim, dim> Avol = make_tensor<dim, dim, dim, dim>(
+          [&](int i, int j, int k, int l)
+          {
+             return (dWvol*J*(invF[j][i]*invF[l][k] - invF[j][k]*invF[l][i])
+                     + J*J*ddWvol*invF[j][i]*invF[l][k]);
+          });
+
+      tensor<double, 3, 3, 3, 3> Adev = make_tensor<dim, dim, dim, dim>(
+          [&](int i, int j, int k, int l)
+          {
+             return 2.0*C1*Jm23*((i==k)*(j==l)
+                                 - 2.0/3.0*(invF[j][i]*F[k][l] + invF[l][k]*F[i][j])
+                                 + I1/3.0*invF[j][k]*invF[l][i] + 2.0/9.0*I1*invF[l][k]*invF[j][i]);
+          });
+      return Adev + Avol;
    }
 
    /**
@@ -120,25 +301,25 @@ struct NeoHookeanMaterial
    action_of_gradient(const tensor<double, dim, dim> &dudx,
                       const tensor<double, dim, dim> &ddudx) const
    {
-      if (gradient_type == GradientType::Symbolic)
+      if (stress_gradient_type == GradientType::Symbolic)
       {
          return action_of_gradient_symbolic(dudx, ddudx);
       }
 #ifdef MFEM_USE_ENZYME
-      else if (gradient_type == GradientType::EnzymeFwd)
+      else if (stress_gradient_type == GradientType::EnzymeFwd)
       {
          return action_of_gradient_enzyme_fwd(dudx, ddudx);
       }
-      else if (gradient_type == GradientType::EnzymeRev)
+      else if (stress_gradient_type == GradientType::EnzymeRev)
       {
          return action_of_gradient_enzyme_rev(dudx, ddudx);
       }
 #endif
-      else if (gradient_type == GradientType::FiniteDiff)
+      else if (stress_gradient_type == GradientType::FiniteDiff)
       {
          return action_of_gradient_fd(dudx, ddudx);
       }
-      else if (gradient_type == GradientType::DualNumbers)
+      else if (stress_gradient_type == GradientType::DualNumbers)
       {
          return action_of_gradient_dual(dudx, ddudx);
       }
@@ -206,20 +387,29 @@ struct NeoHookeanMaterial
    action_of_gradient_symbolic(const tensor<double, dim, dim> &du_dx,
                                const tensor<double, dim, dim> &ddu_dx) const
    {
-      constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      const auto I = mfem::internal::Identity<dim>();
 
-      tensor<double, dim, dim> F = I + du_dx;
-      tensor<double, dim, dim> invFT = inv(transpose(F));
-      tensor<double, dim, dim> devB =
-         dev(du_dx + transpose(du_dx) + dot(du_dx, transpose(du_dx)));
-      double J = det(F);
-      double coef = (C1 / pow(J, 5.0 / 3.0));
-      double a1 = ddot(invFT, ddu_dx);
-      double a2 = ddot(F, ddu_dx);
+      const auto F = I + du_dx;
+      const double J = det(F);
+      const double I1_3 = ddot(F, F)/3.0;
+      const auto invF = inv(F);
+      const auto invFT = transpose(invF);
+      const double fac = pow(J, -2.0/3.0) * 2.0 * C1;
+      const double a1 = ddot(invFT, ddu_dx);
+      const double a2 = ddot(F, ddu_dx);
+      const auto M = dot(invF, dot(ddu_dx, invF));
+      const auto Pdev = fac*(F - I1_3*invFT);
+      
+      auto dPdev = ddu_dx - (2.0/3.0*a2)*invFT + I1_3*transpose(M);
+      dPdev = fac * dPdev;
+      dPdev -= (2.0/3.0*a1) * Pdev;
 
-      return (2.0 * D1 * J * a1 - (4.0 / 3.0) * coef * a2) * I -
-             ((10.0 / 3.0) * coef * a1) * devB +
-             (2 * coef) * (dot(ddu_dx, transpose(F)) + dot(F, transpose(ddu_dx)));
+      const double dWvol = 2.0*D1*(J - 1.0);
+      const double ddWvol = 2.0*D1;
+      auto dPvol = (J*ddWvol + dWvol)*a1*invFT - dWvol*transpose(M);
+      dPvol = J*dPvol;
+      
+      return dPdev + dPvol;
    }
 
    // Parameters

miniapps/elasticity/tests.cpp

@@ -0,0 +1,54 @@
+#include <mfem.hpp>
+
+#include "materials/neohookean.hpp"
+
+using namespace std;
+using namespace mfem;
+
+
+int main()
+{
+  NeoHookeanMaterial<3, GradientType::Symbolic, GradientType::Symbolic> mat;
+  tensor<double, 3, 3> H{{{0.337494265892494, 0.194238454581911, 0.307832573181341},
+                          {0.090147365480304, 0.610402517912401, 0.458978918716148},
+                          {0.689309323130592, 0.198321409053159, 0.901973313462065}}};
+
+  double W = mat.strain_energy_density(H);
+  std::cout << "Strain energy density = " << W << "\n" << endl;
+
+  auto P_symbolic = mat.stress_symbolic(H);
+  auto P_fd = mat.stress_fd(H);
+  auto P_enzyme_rev = mat.stress_enzyme_rev(H);
+  auto P_enzyme_fwd = mat.stress_enzyme_fwd(H);
+  auto P_dual = mat.stress_dual(H);
+  cout << "Stress\n"
+       << "------" << endl;
+  cout << "symbolic\t" << P_symbolic << endl;
+  cout << "finite diff\t" << P_fd << endl;
+  cout << "enzyme rev\t" << P_enzyme_rev << endl;
+  cout << "enzyme fwd\t" << P_enzyme_fwd << endl;
+  cout << "dual numbers\t" << P_dual << endl;
+
+  cout << "\n\n" << endl;
+
+  tensor<double, 3, 3> Hdot{{{0.191653881479253, 0.445956210862074, 0.038732049150475},
+                             {0.589233685844341, 0.092360587237104, 0.259746940075709},
+                             {0.830970655782669, 0.485472875958392, 0.03308538443643 }}};
+
+  auto dP_symbolic = mat.action_of_gradient_symbolic(H, Hdot);
+  auto dP_fd = mat.action_of_gradient_fd(H, Hdot);
+  auto dP_enzyme_rev = mat.action_of_gradient_enzyme_rev(H, Hdot);
+  auto dP_enzyme_fwd = mat.action_of_gradient_enzyme_fwd(H, Hdot);
+  auto dP_dual = mat.action_of_gradient_dual(H, Hdot);
+  auto dP_full = ddot(mat.gradient(H), Hdot);
+  cout << "Elasticities\n"
+       << "------------" << endl;
+  cout << "symbolic\t" << dP_symbolic << endl;
+  cout << "finite diff\t" << dP_fd << endl;
+  cout << "enzyme rev\t" << dP_enzyme_rev << endl;
+  cout << "enzyme fwd\t" << dP_enzyme_fwd << endl;
+  cout << "dual numbers\t" << dP_dual << endl;
+  cout << "contraction of symbolic elastic tensor with \ntangent dispgrad\t" << dP_full << endl;
+
+  return 0;
+}