pMG test

examples/dfem.hpp

@@ -0,0 +1,381 @@
+#pragma once
+
+#include <string_view>
+
+template <typename T>
+constexpr auto get_type_name() -> std::string_view
+{
+#if defined(__clang__)
+   constexpr auto prefix = std::string_view {"[T = "};
+   constexpr auto suffix = "]";
+   constexpr auto function = std::string_view{__PRETTY_FUNCTION__};
+#elif defined(__GNUC__)
+   constexpr auto prefix = std::string_view {"with T = "};
+   constexpr auto suffix = "; ";
+   constexpr auto function = std::string_view{__PRETTY_FUNCTION__};
+#elif defined(_MSC_VER)
+   constexpr auto prefix = std::string_view {"get_type_name<"};
+   constexpr auto suffix = ">(void)";
+   constexpr auto function = std::string_view{__FUNCSIG__};
+#else
+#error Unsupported compiler
+#endif
+
+   const auto start = function.find(prefix) + prefix.size();
+   const auto end = function.find(suffix);
+   const auto size = end - start;
+
+   return function.substr(start, size);
+}
+
+#include <type_traits>
+
+#include <mfem.hpp>
+#include <linalg/dtensor.hpp>
+#include "linalg/tensor.hpp"
+#include <linalg/kernels.hpp>
+
+using namespace mfem;
+
+template <
+   typename type,
+   int rank = 1>
+class NDArray : public Array<type>
+{
+};
+
+template <class F>
+struct FunctionSignature;
+
+template <typename output_type, typename... input_types>
+struct FunctionSignature<output_type(input_types...)>
+{
+   using return_type = output_type;
+   using parameter_types = std::tuple<input_types...>;
+};
+
+template <class T>
+struct create_function_signature;
+
+template <typename output_type, typename T, typename... input_types>
+struct create_function_signature<output_type (T::*)(input_types...) const>
+{
+   using type = FunctionSignature<output_type(input_types...)>;
+};
+
+template <typename function_type, typename... input_types, typename output_type,
+          typename... arg_types>
+void forall_impl(FunctionSignature<output_type(input_types...)>,
+                 const function_type &f, const int n, arg_types &...args)
+{
+   for (int i = 0; i < n; i++)
+   {
+      f(((typename std::remove_reference<input_types>::type
+          *)(args.begin()))[i]...);
+   };
+}
+
+template <typename function_type, typename... arg_types>
+void forall(const function_type &f, const int n, arg_types &...args)
+{
+   using function_signature_type = typename create_function_signature<
+                                   decltype(&function_type::operator())>::type;
+   forall_impl(function_signature_type{}, f, n, args...);
+}
+
+template <typename... arg_types, typename... input_types>
+void forall(void (*f)(arg_types &...), const int n, input_types &...args)
+{
+   for (int i = 0; i < n; i++)
+   {
+      f(((typename std::remove_reference<arg_types>::type*)(args.begin()))[i]...);
+   };
+}
+
+template <
+   typename output_type,
+   typename... input_types>
+auto fwddiff(output_type (*f)(input_types...))
+{
+   return [f](input_types... args, input_types... args2)
+   {
+      auto input_types_tuple = std::tuple<input_types...>(args...);
+      auto shadow_input_types_tuple = std::tuple<input_types...>(args2...);
+
+      static_assert(
+         std::is_same_v<decltype(input_types_tuple), decltype(shadow_input_types_tuple)>,
+         "input and shadow not equal");
+
+      // auto concatenated_types_tuple = std::tuple_cat(input_types_tuple,
+      //                                                shadow_input_types_tuple);
+
+      // std::cout << get_type_name<decltype(concatenated_types_tuple)>() << std::endl;
+
+      return __enzyme_fwddiff<output_type>(f, &args..., &args2...);
+   };
+}
+
+/// @return u_qp qp x vdim x elements
+void interpolate(const GridFunction &u, const IntegrationRule &ir, Vector &u_qp)
+{
+   auto fes = u.FESpace();
+   auto B = fes->GetQuadratureInterpolator(ir);
+   B->SetOutputLayout(QVectorLayout::byVDIM);
+   B->DisableTensorProducts();
+
+   auto R = fes->GetElementRestriction(ElementDofOrdering::NATIVE);
+   Vector u_el(R->Height());
+   R->Mult(u, u_el);
+
+   u_qp.SetSize(fes->GetVDim() *
+                fes->GetMesh()->GetNE() *
+                ir.GetNPoints());
+
+   B->Values(u_el, u_qp);
+}
+
+void gradient_wrt_x(const GridFunction &u, const IntegrationRule &ir,
+                    Vector &grad_u_qp)
+{
+   auto fes = u.FESpace();
+   auto B = fes->GetQuadratureInterpolator(ir);
+   B->SetOutputLayout(QVectorLayout::byVDIM);
+   B->DisableTensorProducts();
+
+   auto R = fes->GetElementRestriction(ElementDofOrdering::NATIVE);
+   Vector u_el(R->Height());
+   R->Mult(u, u_el);
+
+   grad_u_qp.SetSize(
+      fes->GetVDim() *
+      fes->GetMesh()->Dimension() *
+      fes->GetMesh()->GetNE() *
+      ir.GetNPoints());
+
+   B->PhysDerivatives(u_el, grad_u_qp);
+
+   if (fes->GetVDim() > 1)
+   {
+      forall([&](const mfem::internal::tensor<double, 2, 2> &dudx,
+                 mfem::internal::tensor<double, 2, 2> &dudx_transpose)
+      {
+         dudx_transpose = transpose(dudx);
+      }, ir.GetNPoints() * fes->GetMesh()->GetNE(), grad_u_qp, grad_u_qp);
+   }
+}
+
+void integrate_basis(Vector &s_qp, const FiniteElementSpace &fes,
+                     const IntegrationRule &ir, Vector &yi)
+{
+   auto R = fes.GetElementRestriction(ElementDofOrdering::NATIVE);
+
+   auto mesh = fes.GetMesh();
+   // const int dim = mesh->Dimension();
+   const int num_el = mesh->GetNE();
+   const int vdim = fes.GetVDim();
+   const int num_qp = ir.GetNPoints();
+   const int num_vdofs = R->Height() / num_el;
+   const int num_dofs = num_vdofs / vdim;
+
+   if constexpr(false)
+   {
+      out << "#el: " << num_el << " vdim: " << vdim << " #qp: " << num_qp
+          << " #vdofs: " << num_vdofs << " #dofs: " << num_dofs << "\n";
+   }
+
+   const GeometricFactors *geom = mesh->GetGeometricFactors(
+                                     ir, GeometricFactors::JACOBIANS | GeometricFactors::DETERMINANTS);
+   Vector yi_el(R->Height());
+   yi_el = 0.0;
+   auto Yi = Reshape(yi_el.Write(), num_dofs, vdim, num_el);
+   auto C = Reshape(s_qp.ReadWrite(), vdim, num_qp, num_el);
+   auto detJ = Reshape(geom->detJ.Read(), num_qp, num_el);
+   for (int e = 0; e < num_el; e++)
+   {
+      const DofToQuad &maps = fes.GetFE(e)->GetDofToQuad(ir, DofToQuad::FULL);
+      const auto Bt = Reshape(maps.Bt.Read(), num_dofs, num_qp);
+
+      for (int dof = 0; dof < num_dofs; dof++)
+      {
+         for (int vd = 0; vd < vdim; vd++)
+         {
+            double s = 0.0;
+            for (int qp = 0; qp < num_qp; qp++)
+            {
+               s += Bt(dof, qp) * C(vd, qp, e) * detJ(qp, e) * ir.GetWeights()[qp];
+            }
+            Yi(dof, vd, e) = s;
+         }
+      }
+   }
+
+   yi.SetSize(fes.GetVSize());
+   R->MultTranspose(yi_el, yi);
+}
+
+void integrate_basis_gradient(Vector &s_qp, const FiniteElementSpace &fes,
+                              const IntegrationRule &ir, Vector &yi,
+                              const Vector &element_jacobian_inverse)
+{
+   auto R = fes.GetElementRestriction(ElementDofOrdering::NATIVE);
+
+   auto mesh = fes.GetMesh();
+   const int dim = mesh->Dimension();
+   const int num_el = mesh->GetNE();
+   const int vdim = fes.GetVDim();
+   const int num_qp = ir.GetNPoints();
+   const int num_vdofs = R->Height() / num_el;
+   const int num_dofs = num_vdofs / vdim;
+
+   const GeometricFactors *geom = mesh->GetGeometricFactors(
+                                     ir, GeometricFactors::JACOBIANS | GeometricFactors::DETERMINANTS);
+
+   Vector yi_el(R->Height());
+   auto Yi = Reshape(yi_el.Write(), num_dofs, vdim, num_el);
+
+   auto C = Reshape(s_qp.ReadWrite(), vdim, dim, num_qp, num_el);
+   auto detJ = Reshape(geom->detJ.Read(), num_qp, num_el);
+   auto JqpInv = Reshape(element_jacobian_inverse.Read(), num_qp, dim, dim,
+                         num_el);
+
+   CALI_CXX_MARK_LOOP_BEGIN(element_loop, "element_loop");
+   for (int e = 0; e < num_el; e++)
+   {
+      CALI_CXX_MARK_LOOP_ITERATION(element_loop, e);
+      const DofToQuad &maps = fes.GetFE(e)->GetDofToQuad(ir, DofToQuad::FULL);
+      const auto Gt = Reshape(maps.Gt.Read(), num_dofs, num_qp, dim);
+
+      for (int dof = 0; dof < num_dofs; dof++)
+      {
+         for (int vd = 0; vd < vdim; vd++)
+         {
+            double s = 0.0;
+            for (int d = 0; d < dim; d++)
+            {
+               for (int qp = 0; qp < num_qp; qp++)
+               {
+                  const double JxW = detJ(qp, e) * ir.GetWeights()[qp];
+                  for (int k = 0; k < dim; k++)
+                  {
+                     s += Gt(dof, qp, d) * JqpInv(qp, d, k, e) * C(vd, k, qp, e) * JxW;
+                  }
+               }
+            }
+            Yi(dof, vd, e) = s;
+         }
+      }
+
+      // for (int qp = 0; qp < num_qp; qp++)
+      // {
+      //    const double JxW = detJ(qp, e) * ir.GetWeights()[qp];
+
+      //    // Pullback Gradient into physical space
+      //    Vector dphidx_data(num_dofs * dim);
+      //    dphidx_data = 0.0;
+      //    auto dphidx = Reshape(dphidx_data.ReadWrite(), num_dofs, dim);
+      //    for (int d = 0; d < dim; d++)
+      //    {
+      //       for (int dof = 0; dof < num_dofs; dof++)
+      //       {
+      //          for (int k = 0; k < dim; k++)
+      //          {
+      //             dphidx(dof, d) += G(qp, k, dof) * JqpInv(qp, k, d, e);
+      //          }
+      //       }
+      //    }
+
+      //    for (int vd = 0; vd < vdim; vd++)
+      //    {
+      //       for (int dof = 0; dof < num_dofs; dof++)
+      //       {
+      //          double s = 0;
+      //          for (int d = 0; d < dim; d++)
+      //          {
+      //             s += dphidx(dof, d) * C(vd, d, qp, e) * JxW;
+      //          }
+      //          Yi(dof, vd, e) = s;
+      //       }
+      //    }
+      // }
+   }
+   CALI_CXX_MARK_LOOP_END(element_loop);
+   yi.SetSize(fes.GetVSize());
+   R->MultTranspose(yi_el, yi);
+}
+
+void interpolate_boundary(const GridFunction &u, const IntegrationRule &ir_face,
+                          Vector &u_qp)
+{
+   auto fes = u.FESpace();
+   auto B = fes->GetFaceQuadratureInterpolator(ir_face, FaceType::Boundary);
+   B->SetOutputLayout(QVectorLayout::byVDIM);
+   B->DisableTensorProducts();
+
+   auto R = fes->GetFaceRestriction(ElementDofOrdering::LEXICOGRAPHIC,
+                                    FaceType::Boundary);
+   Vector u_el(R->Height());
+   R->Mult(u, u_el);
+
+   u_qp.SetSize(
+      fes->GetVDim() *
+      fes->GetMesh()->GetNBE() *
+      ir_face.GetNPoints());
+
+   B->Values(u_el, u_qp);
+}
+
+void integrate_basis_boundary(Vector &s_qp,
+                              const FiniteElementSpace &fes,
+                              const IntegrationRule &ir_face, Vector &yi)
+{
+   const auto fe = fes.GetFaceElement(0);
+   const auto tfe = dynamic_cast<const TensorBasisElement *>(fe);
+   MFEM_VERIFY(tfe != nullptr, "FE not a TensorBasisElement");
+
+   auto R = fes.GetFaceRestriction(ElementDofOrdering::LEXICOGRAPHIC,
+                                   FaceType::Boundary);
+
+   auto mesh = fes.GetMesh();
+   // const int dim = mesh->Dimension();
+   const int num_fel = mesh->GetNBE();
+   // const int num_fel = mesh->GetNFaces();
+   const int vdim = fes.GetVDim();
+   const int num_qp = ir_face.GetNPoints();
+   const int num_vdofs = R->Height() / num_fel;
+   const int num_dofs = num_vdofs / vdim;
+
+   const FaceGeometricFactors *geom = mesh->GetFaceGeometricFactors(
+                                         ir_face, FaceGeometricFactors::DETERMINANTS,
+                                         FaceType::Boundary, s_qp.GetMemory().GetMemoryType());
+   Vector yi_el(R->Height());
+   yi_el = 0.0;
+   auto Yi = Reshape(yi_el.Write(), num_dofs, vdim, num_fel);
+   auto C = Reshape(s_qp.ReadWrite(), vdim, num_qp, num_fel);
+   auto detJ = Reshape(geom->detJ.Read(), num_qp, num_fel);
+
+   const DofToQuad &maps = fe->GetDofToQuad(ir_face, DofToQuad::FULL);
+   auto lex_to_native = tfe->GetDofMap();
+
+   for (int e = 0; e < num_fel; e++)
+   {
+      // const DofToQuad &maps = fes.GetBE(e)->GetDofToQuad(ir_face, DofToQuad::FULL);
+      const auto Bt = Reshape(maps.Bt.Read(), num_dofs, num_qp);
+
+      for (int dof = 0; dof < num_dofs; dof++)
+      {
+         for (int vd = 0; vd < vdim; vd++)
+         {
+            double s = 0.0;
+            for (int qp = 0; qp < num_qp; qp++)
+            {
+               s += Bt(lex_to_native[dof], qp) * C(vd, qp, e)
+                    * detJ(qp, e) * ir_face.GetWeights()[qp];
+            }
+            Yi(dof, vd, e) = s;
+         }
+      }
+   }
+
+   yi.SetSize(fes.GetVSize());
+   R->MultTranspose(yi_el, yi);
+}

examples/dfem2.cpp

@@ -0,0 +1,666 @@
+#include <cassert>
+#include <functional>
+#include <iostream>
+#include <variant>
+#include <vector>
+
+#include "dfem.hpp"
+
+#include "eigen-3.4.0/Eigen/Eigen"
+
+using namespace mfem;
+using mfem::internal::tensor;
+using mfem::internal::dual;
+
+class LambdaOperator : public Operator
+{
+public:
+    LambdaOperator(int size,
+                   std::function<void(const Vector&, Vector&)> mult_f) :
+        Operator(size),
+        mult_f(mult_f)
+    {}
+
+    void Mult(const Vector& X, Vector& Y) const
+    {
+        mult_f(X, Y);
+    }
+
+    std::function<void(const Vector&, Vector&)> mult_f;
+};
+
+class ADOperator : public Operator
+{
+public:
+    ADOperator(int size) : Operator(size) {}
+    virtual void GradientMult(const Vector &dX, Vector &Y) const = 0;
+    virtual void AdjointMult(const Vector &L, Vector &Y) const = 0;
+};
+
+class PLaplacianGradientOperator : public Operator
+{
+public:
+    PLaplacianGradientOperator(ADOperator &op) :
+        Operator(op.Height()), op(op) {}
+
+    void Mult(const Vector &x, Vector &y) const override
+    {
+        op.GradientMult(x, y);
+    }
+
+    ADOperator &op;
+};
+
+class PLaplacianAdjointOperator : public Operator
+{
+public:
+    PLaplacianAdjointOperator(ADOperator &op) :
+        Operator(op.Height()), op(op) {}
+
+    void Mult(const Vector &x, Vector &y) const override
+    {
+        op.AdjointMult(x, y);
+    }
+
+    ADOperator &op;
+};
+
+class PLaplacianOperator : public ADOperator
+{
+public:
+    PLaplacianOperator(ParFiniteElementSpace &fes, ParGridFunction &u) :
+        ADOperator(fes.GetTrueVSize()),
+        mesh(fes.GetParMesh()),
+        fes(fes),
+        u(u.ParFESpace()),
+        l(u.ParFESpace()),
+        ir(const_cast<IntegrationRule &>(IntRules.Get(
+                                             Element::QUADRILATERAL,
+                                             2 * mesh->GetNodes()->FESpace()->GetOrder(0) + 1)))
+    {
+        Array<int> ess_bdr(mesh->bdr_attributes.Max());
+        ess_bdr = 1;
+        fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+        du.SetSpace(u.FESpace());
+        x_lvec.SetSize(fes.GetProlongationMatrix()->Height());
+        x_lvec = 0.0;
+        y_lvec.SetSize(fes.GetProlongationMatrix()->Height());
+        du_tvec.SetSize(fes.GetProlongationMatrix()->Width());
+        l_tvec.SetSize(fes.GetProlongationMatrix()->Width());
+
+        gradient = new PLaplacianGradientOperator(*this);
+        adjoint = new PLaplacianAdjointOperator(*this);
+    }
+
+    void SetVolumeForcing(ParGridFunction &forcing)
+    {
+        volume_force = &forcing;
+    }
+
+    // F(u, g) = (pow(norm(grad_u), 0.5 * (p - 2)) * grad u, grad phi) - (g, phi)
+    void Mult(const Vector &x, Vector &y) const override
+    {
+        // T -> L
+        fes.GetProlongationMatrix()->Mult(x, u);
+
+        // L -> Q
+
+        // [vdim, dim, num_qp, num_el]
+        auto grad_u_qp = gradient_wrt_x(u, ir);
+
+        // Q -> Q
+        auto plap = [](tensor<double, 2> grad_u)
+        {
+            // grad_u := [dudx dudy]
+            const int p = 4;
+            return (pow(norm(grad_u), 0.5 * (p - 2)) * grad_u);
+        };
+
+        auto f_grad_u_qp = forall(plap, ir.GetNPoints() * mesh->GetNE(),
+                                  grad_u_qp);
+
+        // Layout of f_grad_u_qp_flat has to be [vdim, dim, num_qp, num_el]
+        Vector f_grad_u_qp_flat((double *)f_grad_u_qp.GetData(),
+                                1 * 2 * ir.GetNPoints() *
+                                mesh->GetNE());
+
+        Vector f_grad_u_grad_phi = integrate_basis_gradient(f_grad_u_qp_flat,
+                                   fes,
+                                   ir);
+
+        // - (g, phi)
+        auto g_qp = interpolate(*volume_force, ir);
+        Vector g_phi = integrate_basis(g_qp, fes, ir);
+        f_grad_u_grad_phi -= g_phi;
+
+        // L -> T
+        fes.GetProlongationMatrix()->MultTranspose(f_grad_u_grad_phi, y);
+
+        y.SetSubVector(ess_tdof_list, 0.0);
+    }
+
+    // df/du
+    Operator &GetGradient(const Vector &x) const override
+    {
+        // T -> L
+        fes.GetProlongationMatrix()->Mult(x, state_lvec);
+        return *gradient;
+    }
+
+    // dX: current iterate
+    // Y: dR/dU * dX
+    void GradientMult(const Vector &X, Vector &Y) const override
+    {
+        // apply essential bcs
+        du_tvec = X;
+        du_tvec.SetSubVector(ess_tdof_list, 0.0);
+
+        du.SetFromTrueDofs(du_tvec);
+        u = state_lvec;
+
+        auto grad_u_qp = gradient_wrt_x(u, ir);
+        auto grad_du_qp = gradient_wrt_x(du, ir);
+
+        const int N = mesh->GetNE() * ir.GetNPoints();
+
+        auto plap2 = [](tensor<double, 2> &grad_u)
+        {
+            const int p = 4;
+            return (pow(norm(grad_u), 0.5 * (p - 2)) * grad_u);
+        };
+
+        auto flux_qp = forall([&, plap2](tensor<double,2> grad_u,
+                                         tensor<double,2> grad_du)
+        {
+            return fwddiff(+plap2)(grad_u, grad_du);
+        }, N, grad_u_qp, grad_du_qp);
+
+        // has to be [vdim, dim, num_qp, num_el]
+        Vector flux_qp_flat((double *)flux_qp.GetData(), 1 * 2 * N);
+
+        Vector y = integrate_basis_gradient(flux_qp_flat, fes, ir);
+
+        // L-vector to T-vector
+        fes.GetProlongationMatrix()->MultTranspose(y, Y);
+
+        // Re-assign the essential degrees of freedom on the final output vector.
+        for (int i = 0; i < ess_tdof_list.Size(); i++)
+        {
+            Y[ess_tdof_list[i]] = X[ess_tdof_list[i]];
+        }
+    }
+
+    // (dF/dU)^t
+    Operator &GetAdjoint(const Vector &u) const
+    {
+        // T -> L
+        fes.GetProlongationMatrix()->Mult(u, state_lvec);
+        return *adjoint;
+    }
+
+    // dL: current iterate of adjoint state
+    // Y: (dF/dU)^t * dL
+    void AdjointMult(const Vector &dL, Vector &Y) const override
+    {
+        // apply essential bcs
+        l_tvec = dL;
+        l_tvec.SetSubVector(ess_tdof_list, 0.0);
+
+        l.SetFromTrueDofs(l_tvec);
+        u = state_lvec;
+
+        auto grad_u_qp = gradient_wrt_x(u, ir);
+        // L * B^t -> B * L
+        auto grad_l_qp = gradient_wrt_x(l, ir);
+
+        const int N = mesh->GetNE() * ir.GetNPoints();
+
+        auto plap2 = [](tensor<double, 2> &grad_u, tensor<double, 2> &flux)
+        {
+            const int p = 4;
+            flux = (pow(norm(grad_u), 0.5 * (p - 2)) * grad_u);
+        };
+
+        auto ev_action_qp = forall([&, plap2](tensor<double,2> grad_u,
+                                              tensor<double,2> grad_l)
+        {
+            tensor<double, 2> unused_output{};
+            tensor<double, 2> dgrad_u{};
+            // autodiff == reverse mode
+            __enzyme_autodiff<tensor<double, 2>>(+plap2, &grad_u, &dgrad_u, &unused_output,
+                                                 &grad_l);
+            return dgrad_u;
+        }, N, grad_u_qp, grad_l_qp);
+
+        // has to be [vdim, dim, num_qp, num_el]
+        Vector ev_action_qp_flat((double *)ev_action_qp.GetData(), 1 * 2 * N);
+        Vector y = integrate_basis_gradient(ev_action_qp_flat, fes, ir);
+
+        // L-vector to T-vector
+        fes.GetProlongationMatrix()->MultTranspose(y, Y);
+
+        // Re-assign the essential degrees of freedom on the final output vector.
+        for (int i = 0; i < ess_tdof_list.Size(); i++)
+        {
+            Y[ess_tdof_list[i]] = dL[ess_tdof_list[i]];
+        }
+    }
+
+    // Compute adjoint state of the primal state u
+    Vector ComputeAdjointState(const ParGridFunction &u)
+    {
+        Vector l_tdof(u.ParFESpace()->GetTrueVSize());
+        l_tdof = 0.0; // ?
+        l_tdof.SetSubVector(ess_tdof_list, 0.0);
+
+        // Get adjoint load
+        auto rhs_tdof = ComputeDQoIDU(u);
+        rhs_tdof.SetSubVector(ess_tdof_list, 0.0);
+
+        // Get Jacobian
+        auto u_tdof = u.GetTrueDofs();
+        Operator &J = GetAdjoint(*u_tdof);
+
+        std::ofstream myfile("adjoint.txt");
+        J.PrintMatlab(myfile);
+        myfile.close();
+
+        Operator &G = GetGradient(*u_tdof);
+        std::ofstream myfile2("jacobian.txt");
+        G.PrintMatlab(myfile2);
+        myfile2.close();
+
+        GMRESSolver gmres(MPI_COMM_WORLD);
+        gmres.SetRelTol(1e-12);
+        gmres.SetMaxIter(2000);
+        gmres.SetPrintLevel(0);
+        gmres.SetOperator(J);
+
+        gmres.Mult(rhs_tdof, l_tdof);
+
+        delete u_tdof;
+        return l_tdof;
+    }
+
+    Vector ComputeDfDpTv(const ParGridFunction &v)
+    {
+        const int N = mesh->GetNE() * ir.GetNPoints();
+        Vector dfdpTv(volume_force->ParFESpace()->GetTrueVSize());
+
+        LambdaOperator K(dfdpTv.Size(), [&](const Vector& v, Vector& dfdpTv)
+        {
+            Vector vv(v.Size());
+            // apply essential bcs
+            vv = v;
+            vv.SetSubVector(ess_tdof_list, 0.0);
+
+            GridFunction v_gf(volume_force->FESpace());
+            v_gf.SetFromTrueDofs(vv);
+
+            auto g_qp = interpolate(*volume_force, ir);
+            auto v_qp = interpolate(v_gf, ir);
+
+            auto fg = [](double &g, double &f)
+            {
+                f = -g;
+            };
+
+            auto ev_action_qp = forall([&, fg](double g, double v)
+            {
+                double unused_output = 0.0;
+                double dg = 0.0;
+                __enzyme_autodiff<double>(+fg, &g, &dg, &unused_output, &v);
+                return dg;
+            }, N, g_qp, v_qp);
+
+            Vector ev_action_qp_flat((double *)ev_action_qp.GetData(), N);
+            Vector y = integrate_basis(ev_action_qp_flat, fes, ir);
+
+            // L-vector to T-vector
+            fes.GetProlongationMatrix()->MultTranspose(y, dfdpTv);
+
+            // Re-assign the essential degrees of freedom on the final output vector.
+            for (int i = 0; i < ess_tdof_list.Size(); i++)
+            {
+                dfdpTv[ess_tdof_list[i]] = v[ess_tdof_list[i]];
+            }
+        });
+
+        auto v_tdof = v.GetTrueDofs();
+        K.Mult(*v_tdof, dfdpTv);
+
+        std::ofstream myfile("dfdp.txt");
+        K.PrintMatlab(myfile);
+        myfile.close();
+
+        delete v_tdof;
+        return dfdpTv;
+    }
+
+    double ComputeQoI(const ParGridFunction &u_gf)
+    {
+        const int N = mesh->GetNE() * ir.GetNPoints();
+
+        auto u_qp = interpolate(u_gf, ir);
+
+        auto qoi = [](double u)
+        {
+            return 0.5 * pow(u, 2.0);
+        };
+
+        auto qoi_qp = forall(qoi, 2 * N, u_qp);
+        Vector qoi_qp_flat((double *)qoi_qp.GetData(), 2 * N);
+
+        L2_FECollection l2_0(0, mesh->Dimension());
+        ParFiniteElementSpace l2_0_fes(mesh, &l2_0);
+
+        auto qoi_value = integrate_basis(qoi_qp_flat, l2_0_fes, ir);
+
+        return qoi_value.Sum();
+    }
+
+    Vector ComputeDQoIDU(const ParGridFunction &u_gf)
+    {
+        const int N = mesh->GetNE() * ir.GetNPoints();
+
+        auto u_qp = interpolate(u_gf, ir);
+        Vector du_qp(u_qp);
+        du_qp = 1.0;
+
+        auto qoi = [](double &u)
+        {
+            return 0.5 * pow(u, 2.0);
+        };
+
+        auto dqoidu_qp = forall([&, qoi](double u, double du)
+        {
+            return fwddiff(+qoi)(u, du);
+        }, N, u_qp, du_qp);
+
+        Vector dqoidu_qp_flat((double *)dqoidu_qp.GetData(), N);
+
+        auto dqoidu_qp_value = integrate_basis(dqoidu_qp_flat, *u_gf.FESpace(), ir);
+
+        Vector Y(u_gf.ParFESpace()->GetTrueVSize());
+        u_gf.ParFESpace()->GetProlongationMatrix()->MultTranspose(dqoidu_qp_value, Y);
+
+        return Y;
+    }
+
+    ~PLaplacianOperator()
+    {
+    }
+
+    ParMesh *mesh;
+    ParFiniteElementSpace &fes;
+    mutable ParGridFunction u, du, l, *volume_force = nullptr;
+    IntegrationRule &ir;
+    Array<int> ess_tdof_list;
+    mutable Vector x_lvec, y_lvec, state_lvec, du_tvec, l_tvec;
+
+    PLaplacianGradientOperator *gradient = nullptr;
+    PLaplacianAdjointOperator *adjoint = nullptr;
+
+    bool enable_constant_du = false;
+};
+
+void run_problem6()
+{
+    const int dim = 2;
+
+    Mesh mesh = Mesh::MakeCartesian2D(8, 8, Element::QUADRILATERAL, false,
+                                      2.0*M_PI,
+                                      2.0*M_PI);
+    mesh.EnsureNodes();
+
+    ParMesh pmesh(MPI_COMM_WORLD, mesh);
+
+    auto bdr_attributes = pmesh.bdr_attributes;
+
+    Array<int> ess_attr(bdr_attributes.Max());
+    ess_attr = 1;
+
+    IntegrationRule ir = IntRules.Get(Element::QUADRILATERAL,
+                                      2 * pmesh.GetNodes()->FESpace()->GetOrder(0) + 1);
+
+    H1_FECollection h1fec(2);
+    ParFiniteElementSpace h1fes(&pmesh, &h1fec);
+
+    ParBilinearForm M(&h1fes);
+    auto mass_integrator = new MassIntegrator;
+    mass_integrator->SetIntegrationRule(IntRules.Get(
+                                            Element::QUADRILATERAL,
+                                            2 * pmesh.GetNodes()->FESpace()->GetOrder(0) + 1));
+    M.AddDomainIntegrator(mass_integrator);
+    M.Assemble();
+    M.EliminateEssentialBC(ess_attr);
+    M.Finalize();
+    auto Mmat = M.ParallelAssemble();
+
+    std::ofstream myfile("mass.txt");
+    Mmat->PrintMatlab(myfile);
+    myfile.close();
+
+    ParGridFunction u(&h1fes), g(&h1fes);
+
+    PLaplacianOperator plap(h1fes, u);
+
+    auto coef_g = FunctionCoefficient([](const Vector &x)
+    {
+        return sin(x(0)) * sin(x(1));
+    });
+    g.ProjectCoefficient(coef_g);
+    plap.SetVolumeForcing(g);
+
+    GMRESSolver gmres(MPI_COMM_WORLD);
+    gmres.iterative_mode = false;
+    gmres.SetRelTol(1e-8);
+    gmres.SetMaxIter(10000);
+    gmres.SetPrintLevel(0);
+
+    NewtonSolver newton(MPI_COMM_WORLD);
+    newton.SetPreconditioner(gmres);
+    newton.SetOperator(plap);
+    newton.SetRelTol(1e-8);
+    newton.SetAbsTol(1e-12);
+    newton.SetMaxIter(100);
+    newton.SetPrintLevel(0);
+
+    Vector zero;
+    u.Randomize(1234);
+
+    ConstantCoefficient zero_coeff(0.0);
+    u.ProjectBdrCoefficient(zero_coeff, bdr_attributes);
+
+    Vector *u_tdof = u.GetTrueDofs();
+    newton.Mult(zero, *u_tdof);
+
+    u.SetFromTrueDofs(*u_tdof);
+
+    std::cout << "\nComputing adjoint state\n";
+    auto adjoint_state_tdof = plap.ComputeAdjointState(u);
+
+    ParGridFunction adjoint_state(&h1fes);
+    adjoint_state.SetFromTrueDofs(adjoint_state_tdof);
+
+    Vector dfdpTv = plap.ComputeDfDpTv(adjoint_state);
+
+    Vector dqoidp(dfdpTv);
+    dqoidp.Neg();
+    // adjoint_state.SetFromTrueDofs(dqoidp);
+    // dqoidp.Print();
+
+    // FD test
+    {
+        std::cout << "FD TEST DQoIDu\n";
+
+        auto eval_f = [&](double h)
+        {
+            Vector dqoidu(g.Size());
+            for (int i = 0; i < u.Size(); i++)
+            {
+                // Assign perturbation to input
+                u(i) += h;
+                dqoidu(i) = plap.ComputeQoI(u);
+
+                // Revert perturbation
+                u(i) -= h;
+            }
+
+            return dqoidu;
+        };
+
+        double h = 1e-8;
+        Vector fx = eval_f(0.0);
+        Vector fxph = eval_f(h);
+        fxph -= fx;
+        fxph /= h;
+
+        auto dqoidu = plap.ComputeDQoIDU(u);
+
+        fxph -= dqoidu;
+        std::cout << "|DQoIDU - FD_DQoIDU|_l2 = " << fxph.Norml2() << "\n";
+    }
+
+    // FD test
+    {
+        std::cout << "FD TEST DfDp*v\n";
+
+        double h = 1e-8;
+        Vector fx(dfdpTv), fxph(dfdpTv);
+
+        plap.Mult(*u_tdof, fx);
+
+        adjoint_state *= h;
+        g += adjoint_state;
+        plap.Mult(*u_tdof, fxph);
+        g -= adjoint_state;
+        adjoint_state /= h;
+
+        fxph -= fx;
+        fxph /= h;
+
+        fxph -= dfdpTv;
+        std::cout << "|DfDp*v - FD_DfDp*v|_l2 = " << fxph.Norml2() << "\n";
+    }
+
+    // FD test
+    {
+        std::cout << "FD TEST DQoIDp (total derivative)\n";
+
+        auto eval_f = [&](double h)
+        {
+            Vector dqoidp(g.Size());
+            for (int i = 0; i < u.Size(); i++)
+            {
+                // Assign perturbation to input
+                g(i) += h;
+
+                // u.Randomize(1234);
+                // u.ProjectBdrCoefficient(zero_coeff, ess_attr);
+                // u.GetTrueDofs(*u_tdof);
+                newton.Mult(zero, *u_tdof);
+                u.SetFromTrueDofs(*u_tdof);
+                dqoidp(i) = plap.ComputeQoI(u);
+
+                // Revert perturbation
+                g(i) -= h;
+            }
+
+            return dqoidp;
+        };
+
+        double h = 1e-6;
+        Vector fx = eval_f(0.0);
+        Vector fxph = eval_f(h);
+        fxph -= fx;
+        fxph /= h;
+
+        fxph.SetSubVector(plap.ess_tdof_list, 0.0);
+
+        Vector dqoidp(g.Size());
+        dqoidp = dfdpTv;
+        dqoidp.Neg();
+
+        // fxph.Print(out, fxph.Size());
+        // dqoidp.Print(out, dqoidp.Size());
+
+        fxph -= dqoidp;
+        std::cout << "|DQoIDp - FD_DQoIDp|_l2 = " << fxph.Norml2() << "\n";
+    }
+
+    // auto coef_f = FunctionCoefficient([](const Vector &x)
+    // {
+    //    return 2.0;
+    // });
+    // u.ProjectCoefficient(coef_f);
+
+    // auto qoi = plap.ComputeQoI(u);
+    // std::cout << "QoI = " << qoi << std::endl;
+
+    // auto dqoidu = plap.ComputeDQoIDU(u);
+    // std::cout << "DQoIDU = " << std::endl;
+    // // dqoidu.Print(std::cout, dqoidu.Size());
+
+    // ParLinearForm u_lf(&h1fes);
+    // u_lf.AddDomainIntegrator(new DomainLFIntegrator(coef_f));
+    // u_lf.Assemble();
+    // Vector* u_lf_tdofs = u_lf.ParallelAssemble();
+    // // u_lf_tdofs->Print(std::cout, u_lf_tdofs->Size());
+
+    // *u_lf_tdofs -= dqoidu;
+    // std::cout << "||DQoIDU - EXACT||_L2 = " << u_lf_tdofs->Norml2() << std::endl;
+
+    char vishost[] = "128.15.198.77";
+    int  visport   = 19916;
+    {
+        socketstream sol_sock(vishost, visport);
+        sol_sock << "parallel " << Mpi::WorldSize() << " " << Mpi::WorldRank() << "\n";
+        sol_sock.precision(8);
+        sol_sock << "solution\n" << pmesh << u << std::flush;
+    }
+    {
+        socketstream sol_sock(vishost, visport);
+        sol_sock << "parallel " << Mpi::WorldSize() << " " << Mpi::WorldRank() << "\n";
+        sol_sock.precision(8);
+        sol_sock << "solution\n" << pmesh << adjoint_state << std::flush;
+    }
+}
+
+void run_problem6();
+void run_problem7();
+
+int main(int argc, char *argv[])
+{
+    Mpi::Init();
+    int num_procs = Mpi::WorldSize();
+    int myid = Mpi::WorldRank();
+    Hypre::Init();
+
+    int problem_type = 0;
+
+    OptionsParser args(argc, argv);
+    args.AddOption(&problem_type, "-p", "--problem",
+                   "Problem to run");
+    args.Parse();
+    if (!args.Good())
+    {
+        if (myid == 0)
+        {
+            args.PrintUsage(mfem::out);
+        }
+        return 1;
+    }
+    if (myid == 0)
+    {
+        args.PrintOptions(mfem::out);
+    }
+
+    if (problem_type == 6)
+    {
+        run_problem6();
+    }
+
+    return 0;
+}

examples/dfem_test.cpp

@@ -0,0 +1,350 @@
+#include <functional>
+#include <iostream>
+#include <variant>
+#include <vector>
+
+#include "dfem.hpp"
+
+using namespace mfem;
+using mfem::internal::dual;
+using mfem::internal::tensor;
+
+using namespace std;
+
+int test_integrate_boundary()
+{
+   int polynomial_order = 1;
+
+   Mesh mesh = Mesh::MakeCartesian2D(10, 2, Element::QUADRILATERAL, false, 0.0,
+                                     2.0 * M_PI);
+   mesh.EnsureNodes();
+   ParMesh pmesh(MPI_COMM_WORLD, mesh);
+
+   H1_FECollection h1_fec(polynomial_order);
+   ParFiniteElementSpace h1_fes(&pmesh, &h1_fec);
+
+   auto ir_face = const_cast<IntegrationRule *>(
+                     &IntRules.Get(mesh.GetBdrElementGeometry(0),
+                                   3 * mesh.GetNodes()->FESpace()->GetElementOrder(0) + 1));
+
+   auto h1_prolongation = h1_fes.GetProlongationMatrix();
+
+   ParGridFunction boundary_load(&h1_fes);
+   boundary_load = 0.0;
+
+   VectorFunctionCoefficient boundary_load_coeff(2, [](const Vector &x, Vector &u)
+   {
+      u(0) = 0.0;
+      u(1) = 1.0;
+   });
+
+   {
+      Array<int> boundary_load_attr(pmesh.bdr_attributes.Max());
+      boundary_load_attr = 0;
+      boundary_load_attr[2] = 1;
+      boundary_load.ProjectBdrCoefficient(boundary_load_coeff, boundary_load_attr);
+   }
+
+   Vector boundary_load_qp;
+   interpolate_boundary(boundary_load, *ir_face, boundary_load_qp);
+
+   auto foo = Reshape(boundary_load_qp.Read(), h1_fes.GetVDim(),
+                      ir_face->GetNPoints(), pmesh.GetNBE());
+
+   Vector vec(h1_fes.GetVDim());
+   for (int e = 0; e < pmesh.GetNBE(); e++)
+   {
+      auto Tr = pmesh.GetBdrElementTransformation(e);
+
+      for (int qp = 0; qp < ir_face->GetNPoints(); qp++)
+      {
+         const IntegrationPoint &ip = ir_face->IntPoint(qp);
+
+         Tr->SetIntPoint(&ip);
+
+         boundary_load_coeff.Eval(vec, *Tr, ip);
+
+         out << "(" << ip.x << "," << "y)" << " = " << vec(0) << " " << vec(1) << "\n";
+
+      }
+   }
+
+   return 0;
+}
+
+void compute_element_jacobian_inverse(Mesh &mesh, IntegrationRule *ir,
+                                      Vector &element_jacobian_inverse)
+{
+   const int dim = mesh.Dimension();
+   const int num_el = mesh.GetNE();
+   const int num_qp = ir->GetNPoints();
+
+   element_jacobian_inverse.SetSize(num_qp * dim * dim * num_el);
+
+   // Cache inverse Jacobian on each quadrature point
+   const GeometricFactors *geom = mesh.GetGeometricFactors(
+                                     *ir, GeometricFactors::JACOBIANS);
+   auto J = Reshape(geom->J.Read(), num_qp, dim, dim, num_el);
+   auto Jinv = Reshape(element_jacobian_inverse.Write(), num_qp, dim, dim, num_el);
+   DenseMatrix Jqp(dim, dim), JqpInv(dim, dim);
+   for (int e = 0; e < num_el; e++)
+   {
+      for (int qp = 0; qp < num_qp; qp++)
+      {
+         for (int i = 0; i < dim; i++)
+         {
+            for (int j = 0; j < dim; j++)
+            {
+               Jqp(i, j) = J(qp, i, j, e);
+            }
+         }
+
+         CalcInverse(Jqp, JqpInv);
+
+         for (int i = 0; i < dim; i++)
+         {
+            for (int j = 0; j < dim; j++)
+            {
+               Jinv(qp, i, j, e) = JqpInv(i, j);
+            }
+         }
+      }
+   }
+}
+
+inline
+std::string check_result(double norm, double rtol = 1e-12)
+{
+   if (norm < rtol)
+   {
+      return "✅";
+   }
+   return "❌";
+}
+
+int main(int argc, char *argv[])
+{
+   using namespace std;
+
+   Mpi::Init();
+   int num_procs = Mpi::WorldSize();
+   int myid = Mpi::WorldRank();
+   Hypre::Init();
+
+   int dimension = 2;
+   int polynomial_order = 1;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&polynomial_order, "-o", "--order",
+                  "Finite element order (polynomial degree)");
+   args.ParseCheck();
+
+   std::cout << "Polynomial order = " << polynomial_order << "\n";
+
+   FunctionCoefficient linear_scalar_coeff([&](const Vector &c)
+   {
+      double x = c(0), y = c(1);
+      return 2.0 * x + x * y;
+   });
+
+   VectorFunctionCoefficient dlinear_scalardx_coeff(dimension, [&](const Vector &c,
+                                                                   Vector &u)
+   {
+      double x = c(0), y = c(1);
+      u(0) = 2.0 + c(1);
+      u(1) = c(0);
+   });
+
+   FunctionCoefficient quadratic_coeff([&](const Vector &c)
+   {
+      double x = c(0), y = c(1);
+      return 2.0*x*x + x*y*y;
+   });
+
+   VectorFunctionCoefficient dquadraticdx_coeff(dimension, [&](const Vector &c,
+                                                               Vector &u)
+   {
+      double x = c(0), y = c(1);
+      u(0) = 4.0*x+y*y,
+      u(1) = 2.0*x*y;
+   });
+
+   {
+      Mesh mesh = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL, false, 1.0,
+                                        1.0);
+      mesh.EnsureNodes();
+      ParMesh pmesh(MPI_COMM_WORLD, mesh);
+
+      H1_FECollection h1_fec(polynomial_order);
+      ParFiniteElementSpace h1_fes(&pmesh, &h1_fec);
+      ParFiniteElementSpace h1_vfes(&pmesh, &h1_fec, dimension, Ordering::byVDIM);
+
+      cout << "#dofs: " << h1_fes.GetVSize() << "\n\n";
+
+      auto ir = const_cast<IntegrationRule *>(
+                   &IntRules.Get(mesh.GetElementGeometry(0),
+                                 2 * mesh.GetNodes()->FESpace()->GetElementOrder(0)));
+
+      Vector element_jacobian_inverse;
+      compute_element_jacobian_inverse(mesh, ir, element_jacobian_inverse);
+
+      auto h1v_prolongation = h1_vfes.GetProlongationMatrix();
+
+      ParGridFunction u(&h1_fes), du(&h1_vfes), uv(&h1_vfes);
+      u = 0.0, du = 0.0, uv = 0.0;
+
+      {
+         cout << "scalar interpolation\n";
+         Vector u_qp;
+         u.ProjectCoefficient(linear_scalar_coeff);
+         interpolate(u, *ir, u_qp);
+         integrate_basis(u_qp, h1_fes, *ir, u);
+         double integral = 0.0;
+         for (int dof = 0; dof < h1_fes.GetVSize(); dof++)
+         {
+            integral += u(dof);
+         }
+         cout << "|I[u]dx - I[u_ex]dx| = " << abs(integral - 5.0/4.0) << "\n";
+         cout << endl;
+      }
+
+      {
+         cout << "weak gradient of scalar\n";
+         Vector dudx_qp;
+         u.ProjectCoefficient(linear_scalar_coeff);
+         gradient_wrt_x(u, *ir, dudx_qp);
+         integrate_basis(dudx_qp, h1_vfes, *ir, du);
+
+         Vector integral(2);
+         for (int d = 0; d < du.FESpace()->GetVDim(); d++)
+         {
+            integral(d) = 0.0;
+            for (int i = 0; i < du.FESpace()->GetNDofs(); i++)
+            {
+               int idx = Ordering::Map<Ordering::byVDIM>(
+                            du.FESpace()->GetNDofs(),
+                            du.FESpace()->GetVDim(),
+                            i,
+                            d);
+               integral(d) += du(idx);
+            }
+         }
+         cout << "|I[du]dx - I[du_ex]dx| = " << abs(integral(0) - 5.0/2.0) << "\n"
+              << "|I[du]dy - I[du_ex]dy| = " << abs(integral(1) - 1.0/2.0) << "\n";
+
+         ParLinearForm l(&h1_vfes);
+         auto integrator = new VectorDomainLFIntegrator(dlinear_scalardx_coeff);
+         integrator->SetIntRule(ir);
+         l.AddDomainIntegrator(integrator);
+         l.Assemble();
+         du -= *l.ParallelAssemble();
+         cout << "|du - du_form|_l2 = " << du.Norml2()
+              << check_result(du.Norml2()) << "\n";
+
+         cout << endl;
+      }
+
+      {
+         cout << "scalar diffusion, linear u\n";
+         Vector dudx_qp, ru(h1_fes.GetVSize());
+         u.ProjectCoefficient(linear_scalar_coeff);
+         gradient_wrt_x(u, *ir, dudx_qp);
+         integrate_basis_gradient(dudx_qp, h1_fes, *ir, ru,
+                                  element_jacobian_inverse);
+
+         ParBilinearForm b(&h1_fes);
+         auto integrator = new DiffusionIntegrator;
+         integrator->SetIntRule(ir);
+         b.AddDomainIntegrator(integrator);
+         b.Assemble();
+         b.Finalize();
+
+         ParGridFunction y(&h1_fes);
+         b.Mult(u, y);
+
+         y -= ru;
+         cout << "|r(u) - r(u)_form|_l2 = " << y.Norml2()
+              << check_result(y.Norml2()) << "\n";
+
+         cout << endl;
+      }
+
+      {
+         cout << "scalar diffusion, quadratic u\n";
+         Vector dudx_qp, ru(h1_fes.GetVSize());
+         u.ProjectCoefficient(quadratic_coeff);
+         gradient_wrt_x(u, *ir, dudx_qp);
+         integrate_basis_gradient(dudx_qp, h1_fes, *ir, ru,
+                                  element_jacobian_inverse);
+
+         ParBilinearForm b(&h1_fes);
+         auto integrator = new DiffusionIntegrator;
+         integrator->SetIntRule(ir);
+         b.AddDomainIntegrator(integrator);
+         b.Assemble();
+         b.Finalize();
+
+         ParGridFunction y(&h1_fes);
+         b.Mult(u, y);
+
+         y -= ru;
+         cout << "|r(u) - r(u)_form|_l2 = " << y.Norml2()
+              << check_result(y.Norml2()) << "\n";
+
+         cout << endl;
+      }
+
+      {
+         cout << "vector diffusion, linear u\n";
+         Vector duvdx_qp, ru(h1_vfes.GetVSize());
+         uv.ProjectCoefficient(dlinear_scalardx_coeff);
+         gradient_wrt_x(uv, *ir, duvdx_qp);
+         integrate_basis_gradient(duvdx_qp, h1_vfes, *ir, ru,
+                                  element_jacobian_inverse);
+
+         ParBilinearForm b(&h1_vfes);
+         auto integrator = new VectorDiffusionIntegrator;
+         integrator->SetIntRule(ir);
+         b.AddDomainIntegrator(integrator);
+         b.Assemble();
+         b.Finalize();
+
+         ParGridFunction y(&h1_vfes);
+         b.Mult(uv, y);
+
+         y -= ru;
+         cout << "|r(u) - r(u)_form|_l2 = " << y.Norml2()
+              << check_result(y.Norml2()) << "\n";
+
+         cout << endl;
+      }
+
+      {
+         cout << "vector diffusion, quadratic u\n";
+         Vector duvdx_qp, ru(h1_vfes.GetVSize());
+         uv.ProjectCoefficient(dquadraticdx_coeff);
+         gradient_wrt_x(uv, *ir, duvdx_qp);
+         integrate_basis_gradient(duvdx_qp, h1_vfes, *ir, ru,
+                                  element_jacobian_inverse);
+
+         ParBilinearForm b(&h1_vfes);
+         auto integrator = new VectorDiffusionIntegrator;
+         integrator->SetIntRule(ir);
+         b.AddDomainIntegrator(integrator);
+         b.Assemble();
+         b.Finalize();
+
+         ParGridFunction y(&h1_vfes);
+         b.Mult(uv, y);
+
+         y -= ru;
+         cout << "|r(u) - r(u)_form|_l2 = " << y.Norml2()
+              << check_result(y.Norml2()) << "\n";
+
+         cout << endl;
+      }
+
+   }
+
+   return 0;
+}

examples/hooke2.cpp

@@ -0,0 +1,900 @@
+#include <functional>
+#include <iostream>
+#include <variant>
+#include <vector>
+
+#include "dfem.hpp"
+
+#include <caliper/cali.h>
+#include <caliper/cali-manager.h>
+
+using namespace mfem;
+using mfem::internal::tensor;
+using mfem::internal::dual;
+using mfem::internal::make_tensor;
+
+template <int dim>
+class AffineSolution
+{
+public:
+   AffineSolution() : A(dim), b(dim)
+   {
+      // clang-format off
+      A(0, 0) = 0.110791568544027; A(0, 1) = 0.230421268325901;
+      A(1, 0) = 0.198344644470483; A(1, 1) = 0.060514559793513;
+      if constexpr (dim == 3)
+      {
+         A(0, 2) = 0.15167673653354;
+         A(1, 2) = 0.084137393813728;
+         A(2, 0) = 0.011544253485023; A(2, 1) = 0.060942846497753;
+         A(2, 2) = 0.186383473579596;
+      }
+      A *= 1e-2;
+
+      b(0) = 0.765645367640828;
+      b(1) = 0.992487355850465;
+      if constexpr (dim == 3)
+      {
+         b(2) = 0.162199373722092;
+      }
+      b *= 1e-2;
+      //clang-format on
+   };
+
+   /**
+    * @brief MFEM-style coefficient function corresponding to this solution
+    *
+    * @param X Coordinates of point in reference configuration at which solution is sought
+    * @param u Exact solution evaluated at \p X
+    */
+   void operator()(const mfem::Vector& X, mfem::Vector& u) const
+   {
+      A.Mult(X, u);
+      u += b;
+   }
+
+   // /**
+   //  * @brief Apply forcing that should produce this exact displacement
+   //  *
+   //  * Given the physics module, apply boundary conditions and a source
+   //  * term that are consistent with the exact solution. This is
+   //  * independent of the domain. The solution is imposed as an essential
+   //  * boundary condition on the parts of the boundary identified by \p
+   //  * essential_boundaries. On the complement of
+   //  * \p essential_boundaries, the traction corresponding to the exact
+   //  * solution is applied.
+   //  *
+   //  * @tparam p Polynomial degree of the finite element approximation
+   //  * @tparam Material Type of the material model used in the problem
+   //  *
+   //  * @param material Material model used in the problem
+   //  * @param sf The SolidMechanics module for the problem
+   //  * @param essential_boundaries Boundary attributes on which essential boundary conditions are desired
+   //  */
+   // template <int p, typename Material>
+   // void applyLoads(const Material& material, SolidMechanics<p, dim>& sf,
+   //                 std::set<int> essential_boundaries) const
+   // {
+   //    // essential BCs
+   //    auto ebc_func = [*this](const auto& X, auto& u) { this->operator()(X, u); };
+   //    sf.setDisplacementBCs(essential_boundaries, ebc_func);
+
+   //    // natural BCs
+   //    typename Material::State state;
+   //    auto H = make_tensor<dim, dim>([&](int i, int j) { return A(i,j); });
+   //    tensor<double, dim, dim> sigma = material(state, H);
+   //    auto P = solid_mechanics::CauchyToPiola(sigma, H);
+   //    auto traction = [P](auto, auto n0, auto) { return dot(P, n0); };
+   //    sf.setPiolaTraction(traction);
+   // }
+
+private:
+   /// Linear part of solution. Equivalently, the displacement gradient
+   mfem::DenseMatrix A;
+   /// Constant part of solution. Rigid mody displacement.
+   mfem::Vector b;
+};
+
+
+class LambdaOperator : public Operator
+{
+public:
+   LambdaOperator(int size,
+                  std::function<void(const Vector&, Vector&)> mult_f) :
+      Operator(size),
+      mult_f(mult_f)
+   {}
+
+   void Mult(const Vector& X, Vector& Y) const
+   {
+      mult_f(X, Y);
+   }
+
+   std::function<void(const Vector&, Vector&)> mult_f;
+};
+
+class ADOperator : public Operator
+{
+public:
+   ADOperator(int size = 0) : Operator(size) {}
+   virtual void GradientMult(const Vector &dX, Vector &Y) const = 0;
+   virtual void AdjointMult(const Vector &L, Vector &Y) const = 0;
+};
+
+class ElasticityGradientOperator : public Operator
+{
+public:
+   ElasticityGradientOperator(ADOperator &op) :
+      Operator(op.Height()), op(op) {}
+
+   void Mult(const Vector &x, Vector &y) const override
+   {
+      op.GradientMult(x, y);
+   }
+
+   ADOperator &op;
+};
+
+void finite_stress_qf(const tensor<double, 2, 2> &dudx, tensor<double, 2, 2> &P)
+{
+   double lambda, mu;
+   {
+      lambda = 1.25;
+      mu = 1.0;
+   }
+   static constexpr auto I = mfem::internal::IsotropicIdentity<2>();
+   auto F = dudx + I;
+   auto E = 0.5 * (transpose(F) * F - I);
+   // auto eps = sym(dudx);
+   // auto dudx_squared = transpose(dudx) * dudx;
+   // auto E = eps + 0.5 * dudx_squared;
+   auto S = lambda * tr(E) * I + 2.0 * mu * E;
+   P = F * S;
+};
+
+// linear elastic
+void small_stress_qf(const tensor<double, 2, 2> &dudx,
+                     tensor<double, 2, 2> &P)
+{
+   double lambda, mu;
+   {
+      lambda = 1.25;
+      mu = 1.0;
+   }
+   static constexpr auto I = mfem::internal::IsotropicIdentity<2>();
+   auto eps = sym(dudx);
+   auto S = lambda * tr(eps) * I + 2.0 * mu * eps;
+   P = S;
+};
+
+template <auto quadrature_function>
+class ElasticityOperator : public ADOperator
+{
+public:
+   ElasticityOperator(ParMesh &mesh, ParFiniteElementSpace &h1_fes, bool matfree,
+                      bool dump_matrices) :
+      ADOperator(),
+      mesh(mesh),
+      dim(mesh.Dimension()),
+      vdim(mesh.Dimension()),
+      num_el(mesh.GetNE()),
+      h1_fes(h1_fes),
+      matfree(matfree),
+      dump_matrices(dump_matrices)
+   {
+      this->height = h1_fes.GetTrueVSize();
+      this->width = this->height;
+
+      ir = const_cast<IntegrationRule *>(
+              &IntRules.Get(mesh.GetElementGeometry(0),
+                            2 * h1_fes.GetElementOrder(0)));
+
+      ir_face = const_cast<IntegrationRule *>(
+                   &IntRules.Get(mesh.GetBdrElementGeometry(0),
+                                 2 * h1_fes.GetElementOrder(0)));
+
+      num_qp = ir->GetNPoints();
+
+      int global_tdof_size = h1_fes.GlobalTrueVSize();
+      if (Mpi::Root())
+      {
+         out << "dim = " << mesh.Dimension() << "\n"
+             << "vdim = " << h1_fes.GetVDim() << "\n"
+             << "#dofs: " << global_tdof_size << "\n"
+             << "#qp in IntRule: " << num_qp << std::endl;
+      }
+
+      h1_prolongation = h1_fes.GetProlongationMatrix();
+
+      u.SetSpace(&h1_fes);
+      current_state.SetSpace(&h1_fes);
+      current_iterate.SetSpace(&h1_fes);
+      current_iterate_tvec.SetSize(h1_prolongation->Width());
+
+      body_force.SetSpace(&h1_fes);
+      body_force = 0.0;
+
+      boundary_load.SetSpace(&h1_fes);
+      boundary_load = 0.0;
+
+      // Layout has to be [vdim, dim, num_qp, num_el]
+      P_dudx_qp.SetSize(dim * dim * num_qp * num_el);
+      out_qp.SetSize(dim * dim * num_qp * num_el);
+
+      element_jacobian_inverse.SetSize(num_qp * dim * dim * num_el);
+
+      matfree_gradient = new ElasticityGradientOperator(*this);
+
+      // Cache inverse Jacobian on each quadrature point
+      {
+         const GeometricFactors *geom = mesh.GetGeometricFactors(
+                                           *ir, GeometricFactors::JACOBIANS);
+         auto J = Reshape(geom->J.Read(), num_qp, dim, dim, num_el);
+         auto Jinv = Reshape(element_jacobian_inverse.Write(), num_qp, dim, dim, num_el);
+         DenseMatrix Jqp(dim, dim), JqpInv(dim, dim);
+         for (int e = 0; e < num_el; e++)
+         {
+            for (int qp = 0; qp < num_qp; qp++)
+            {
+               for (int i = 0; i < dim; i++)
+               {
+                  for (int j = 0; j < dim; j++)
+                  {
+                     Jqp(i, j) = J(qp, i, j, e);
+                  }
+               }
+
+               CalcInverse(Jqp, JqpInv);
+
+               for (int i = 0; i < dim; i++)
+               {
+                  for (int j = 0; j < dim; j++)
+                  {
+                     Jinv(qp, i, j, e) = JqpInv(i, j);
+                  }
+               }
+            }
+         }
+      }
+   }
+
+   void Mult(const Vector &X, Vector &Y) const override
+   {
+      CALI_MARK_BEGIN("ElasticityOperator::Mult");
+
+      h1_prolongation->Mult(X, u);
+
+      gradient_wrt_x(u, *ir, dudx_qp);
+
+      forall(quadrature_function, num_qp * num_el, dudx_qp, P_dudx_qp);
+
+      integrate_basis_gradient(P_dudx_qp,
+                               h1_fes,
+                               *ir, P_dudx_qp_dphi,
+                               element_jacobian_inverse);
+
+      if (enable_body_force)
+      {
+         interpolate(body_force, *ir, body_force_qp);
+         integrate_basis(body_force_qp, h1_fes, *ir, body_force_phi);
+         P_dudx_qp_dphi += body_force_phi;
+      }
+      if (enable_boundary_load)
+      {
+         interpolate_boundary(boundary_load, *ir_face, boundary_load_qp);
+         integrate_basis_boundary(boundary_load_qp, h1_fes, *ir_face,
+                                  boundary_load_phi);
+         P_dudx_qp_dphi -= boundary_load_phi;
+      }
+
+      h1_prolongation->MultTranspose(P_dudx_qp_dphi, Y);
+      Y.SetSubVector(ess_tdof_list, 0.0);
+
+      CALI_MARK_END("ElasticityOperator::Mult");
+   }
+
+   Operator &GetGradientNoBC(const Vector &x) const
+   {
+      assemble_with_bc = false;
+      auto& op = GetGradient(x);
+      assemble_with_bc = true;
+      return op;
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      // T -> L
+      h1_fes.GetProlongationMatrix()->Mult(x, current_state);
+
+      // Cache dudx
+      gradient_wrt_x(current_state, *ir, dudx_qp);
+
+      if (!matfree)
+      {
+         if (dump_matrices)
+         {
+            std::ofstream matfree_gradient_out("matfreeA.txt");
+            matfree_gradient->PrintMatlab(matfree_gradient_out);
+            matfree_gradient_out.close();
+         }
+
+         Vector dPddudx_qp(dim * dim * dim * dim * num_qp * num_el);
+         CALI_MARK_BEGIN("EnzymeAD Jacobian Assemble");
+         forall([&](const tensor<double, 2, 2> &dudx,
+                    tensor<double, 2, 2, 2, 2> &dPddudx)
+         {
+            tensor<double, 2, 2> unused_output{};
+            tensor<double, 2, 2> dir{};
+            dPddudx = {};
+            for (int i = 0; i < dim; i++)
+            {
+               for (int j = 0; j < dim; j++)
+               {
+                  dir[i][j] = 1;
+                  __enzyme_autodiff<void>(+quadrature_function,
+                                          enzyme_dup, &dudx, &(dPddudx[j][i]), // autodiff returns A^t
+                                          enzyme_dupnoneed, &unused_output, &dir);
+                  dir[i][j] = 0;
+               }
+            }
+         }, num_qp * num_el, dudx_qp, dPddudx_qp);
+         CALI_MARK_END("EnzymeAD Jacobian Assemble");
+
+         // Assemble processor local SparseMatrix
+         SparseMatrix *mat = new SparseMatrix(h1_fes.GetVSize());
+         {
+            auto R = h1_fes.GetElementRestriction(ElementDofOrdering::NATIVE);
+            const int dim = mesh.Dimension();
+            const int num_el = mesh.GetNE();
+            const int vdim = h1_fes.GetVDim();
+            const int num_qp = ir->GetNPoints();
+            const int num_vdofs = R->Height() / num_el;
+            const int num_dofs = num_vdofs / vdim;
+
+            const GeometricFactors *geom = mesh.GetGeometricFactors(
+                                              *ir, GeometricFactors::JACOBIANS | GeometricFactors::DETERMINANTS);
+            auto detJ = Reshape(geom->detJ.Read(), num_qp, num_el);
+            auto invJ = Reshape(element_jacobian_inverse.Read(), num_qp, dim, dim, num_el);
+
+            Vector A_l(num_dofs * dim * num_dofs * dim * num_el);
+            A_l = 0.0;
+
+            auto A_e = Reshape(A_l.ReadWrite(), num_dofs, dim, num_dofs, dim, num_el);
+            auto D = Reshape(dPddudx_qp.Read(), dim, dim, dim, dim, num_qp, num_el);
+
+            for (int e = 0; e < num_el; e++)
+            {
+               const DofToQuad &maps = h1_fes.GetFE(e)->GetDofToQuad(*ir, DofToQuad::FULL);
+               const auto G = Reshape(maps.G.Read(), num_qp, dim, num_dofs);
+
+               for (int qp = 0; qp < num_qp; qp++)
+               {
+                  const double JxW = detJ(qp, e) * ir->GetWeights()[qp];
+
+                  // Pullback Gradient into physical space
+                  Vector dphidx_data(num_dofs * dim);
+                  dphidx_data = 0.0;
+                  auto dphidx = Reshape(dphidx_data.ReadWrite(), num_dofs, dim);
+                  for (int j = 0; j < num_dofs; j++)
+                  {
+                     for (int i = 0; i < dim; i++)
+                     {
+                        double s = 0.0;
+                        for (int k = 0; k < dim; k++)
+                        {
+                           s += G(qp, k, j) * invJ(qp, k, i, e);
+                        }
+                        dphidx(j, i) += s;
+                     }
+                  }
+
+                  for (int q = 0; q < dim; q++)
+                  {
+                     for (int b = 0; b < num_dofs; b++)
+                     {
+                        for (int i = 0; i < dim; i++)
+                        {
+                           for (int a = 0; a < num_dofs; a++)
+                           {
+                              double s = 0.0;
+                              for (int l = 0; l < dim; l++)
+                              {
+                                 for (int k = 0; k < dim; k++)
+                                 {
+                                    // dN^A/dX_k (D_ijkl dN^C/dX_l)
+                                    s += dphidx(a,k) * D(i,k,q,l,qp,e) * dphidx(b,l);
+
+                                    // diagonal test
+                                    // s += dphidx(a,k) * D(i,k,q,l,qp,e) * dphidx(a,l);
+                                 }
+                              }
+                              A_e(a, i, b, q, e) += s * JxW;
+
+                              // diagonal test
+                              // A_e(a, i, a, q, e) += s * JxW;
+                           }
+                        }
+                     }
+                  }
+               }
+            }
+
+            for (int e = 0; e < num_el; e++)
+            {
+               auto tmp = Reshape(A_l.ReadWrite(), num_dofs, dim, num_dofs, dim, num_el);
+               DenseMatrix A_e(&tmp(0, 0, 0, 0, e), num_vdofs, num_vdofs);
+               Array<int> vdofs;
+               h1_fes.GetElementVDofs(e, vdofs);
+               mat->AddSubMatrix(vdofs, vdofs, A_e, 1);
+            }
+            mat->Finalize();
+
+            auto tmp = new HypreParMatrix(h1_fes.GetComm(),
+                                          h1_fes.GlobalVSize(),
+                                          h1_fes.GetDofOffsets(),
+                                          mat);
+            delete Amat;
+            Amat = RAP(tmp, h1_fes.Dof_TrueDof_Matrix());
+            delete tmp;
+            delete mat;
+
+            if (assemble_with_bc)
+            {
+               Amat->EliminateBC(ess_tdof_list, DiagonalPolicy::DIAG_ONE);
+            }
+
+            if (dump_matrices)
+            {
+               std::ofstream assembled_jacobian_out("assembled_jacobian.txt");
+               Amat->PrintMatlab(assembled_jacobian_out);
+               assembled_jacobian_out.close();
+               out << "exiting after writing jacobian matrices to disk...\n";
+               exit(0);
+            }
+
+            return *Amat;
+         }
+      }
+      else
+      {
+         return *matfree_gradient;
+      }
+   }
+
+   // X: current iterate
+   // Y: dR/dU * X
+   void GradientMult(const Vector &X, Vector &Y) const override
+   {
+      CALI_MARK_BEGIN("ElasticityOperator::GradientMult");
+
+      // apply essential bcs
+      current_iterate_tvec = X;
+      current_iterate_tvec.SetSubVector(ess_tdof_list, 0.0);
+      current_iterate.SetFromTrueDofs(current_iterate_tvec);
+
+      gradient_wrt_x(current_iterate, *ir, ddudx_qp);
+
+      CALI_MARK_BEGIN("EnzymeAD MatVec");
+      forall([](const tensor<double, 2, 2> &dudx,
+                tensor<double, 2, 2> &ddudx,
+                tensor<double, 2, 2> &out)
+      {
+         tensor<double, 2, 2> unused_output{};
+         out = {};
+         __enzyme_fwddiff<void>(+quadrature_function, &dudx, &ddudx,
+                                &unused_output, &out);
+      }, num_qp * num_el, dudx_qp, ddudx_qp, out_qp);
+      CALI_MARK_END("EnzymeAD MatVec");
+
+      CALI_MARK_BEGIN("integrate_basis_gradient");
+      integrate_basis_gradient(out_qp, h1_fes, *ir, y, element_jacobian_inverse);
+      CALI_MARK_END("integrate_basis_gradient");
+
+      // L-vector to T-vector
+      h1_fes.GetProlongationMatrix()->MultTranspose(y, Y);
+
+      // Re-assign the essential degrees of freedom on the final output vector.
+      for (int i = 0; i < ess_tdof_list.Size(); i++)
+      {
+         Y[ess_tdof_list[i]] = X[ess_tdof_list[i]];
+      }
+
+      CALI_MARK_END("ElasticityOperator::GradientMult");
+   }
+
+   // dL: current iterate of adjoint state
+   // Y: (dF/dU)^t * dL
+   void AdjointMult(const Vector &dL, Vector &Y) const override
+   {
+
+   }
+
+   void SetEssentialAttributes(const Array<int> attr)
+   {
+      h1_fes.GetEssentialTrueDofs(attr, ess_tdof_list);
+   }
+
+   void SetPrescribedDisplacement(const Array<int> attr)
+   {
+      h1_fes.GetEssentialTrueDofs(attr, displaced_tdof_list);
+   }
+
+   const Array<int> &GetPrescribedDisplacementTDofs()
+   {
+      return displaced_tdof_list;
+   };
+
+   ParGridFunction* GetExternalLoad()
+   {
+      enable_boundary_load = true;
+      return &boundary_load;
+   }
+
+   ParGridFunction* GetBodyForce()
+   {
+      enable_body_force = true;
+      return &body_force;
+   }
+
+   ParMesh &mesh;
+   const int dim;
+   const int vdim;
+   /// Number of elements in the mesh (rank local)
+   int num_el;
+   int num_qp = 0;
+   /// H1 finite element space
+   ParFiniteElementSpace &h1_fes;
+   // Integration rule
+   IntegrationRule *ir = nullptr, *ir_face = nullptr;
+   const Operator *h1_element_restriction;
+   const Operator *h1_prolongation;
+
+   mutable Vector element_jacobian_inverse;
+
+   Array<int> ess_tdof_list, displaced_tdof_list;
+
+   ParGridFunction body_force, boundary_load;
+   mutable ParGridFunction u, current_state, current_iterate;
+   mutable Vector current_iterate_tvec, dudx_qp, ddudx_qp,
+           P_dudx_qp_dphi, body_force_qp, boundary_load_qp, body_force_phi,
+           boundary_load_phi, y,
+           P_dudx_qp, out_qp;
+
+   bool enable_boundary_load = false, enable_body_force = false;
+   ElasticityGradientOperator *matfree_gradient = nullptr;
+
+   bool matfree;
+   bool dump_matrices;
+   mutable HypreParMatrix *Amat = nullptr;
+   mutable bool assemble_with_bc = true;
+};
+
+int main(int argc, char *argv[])
+{
+   using namespace std;
+
+   Mpi::Init();
+   // int num_procs = Mpi::WorldSize();
+   // int myid = Mpi::WorldRank();
+   Hypre::Init();
+
+   bool visualization = false;
+   int polynomial_degree = 1;
+   int refinements = 0;
+   int problem_type = 0;
+   bool matfree = false;
+   bool pmg = false;
+   bool dump_matrices = false;
+   const char *caliper_options = "";
+   const char *mesh_file = "../data/beam-quad.mesh";
+
+   std::shared_ptr<VectorFunctionCoefficient> u_ex_coeff;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh",
+                  "Mesh file to use.");
+   args.AddOption(&polynomial_degree, "-o", "--order",
+                  "Finite element order (polynomial degree)");
+   args.AddOption(&refinements, "-r", "--ref",
+                  "");
+   args.AddOption(&problem_type, "-p", "--problem",
+                  "");
+   args.AddOption(&visualization, "-vis", "--visualization", "-no-vis",
+                  "--no-visualization",
+                  "Enable or disable GLVis visualization.");
+   args.AddOption(&matfree, "-matfree", "--matfree", "-no-matfree",
+                  "--no-matfree",
+                  "matrix free");
+   args.AddOption(&pmg, "-pmg", "--pmg", "-no-pmg",
+                  "--no-pmg",
+                  "p-Multigrid");
+   args.AddOption(&dump_matrices, "-dump_matrices", "--dump_matrices",
+                  "-no-dump_matrices",
+                  "--no-dump_matrices",
+                  "dump matrices");
+   args.AddOption(&caliper_options, "-profile", "--profile", "caliper options");
+   args.ParseCheck();
+
+   if (problem_type == 3)
+   {
+      // MFEM_ASSERT(strcmp("patch2D_quads.mesh", mesh_file) == 0,
+      //             "have to use patch2D_quads.mesh");
+   }
+
+   cali::ConfigManager caliper_mgr;
+   caliper_mgr.add(caliper_options);
+
+   caliper_mgr.start();
+   CALI_MARK_FUNCTION_BEGIN;
+
+   Mesh mesh(mesh_file, 1, 1);
+   mesh.EnsureNodes();
+   int dim = mesh.Dimension();
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh.UniformRefinement();
+   }
+
+   ParMesh pmesh(MPI_COMM_WORLD, mesh);
+   int num_el_global = pmesh.GetGlobalNE();
+   if (Mpi::Root())
+   {
+      out << "#elements: " << num_el_global << "\n";
+   }
+
+   auto *fec = new H1_FECollection(1, dim);
+   auto *coarse_fespace = new ParFiniteElementSpace(&pmesh, fec, dim);
+
+   Array<FiniteElementCollection*> collections;
+   collections.Append(fec);
+   auto* fespaces = new ParFiniteElementSpaceHierarchy(&pmesh, coarse_fespace,
+                                                       true, true);
+   for (int level = 0; level < polynomial_degree; ++level)
+   {
+      collections.Append(new H1_FECollection((int)std::pow(2, level+1), dim));
+      fespaces->AddOrderRefinedLevel(collections.Last(), dim);
+   }
+
+   HYPRE_BigInt size = fespaces->GetFinestFESpace().GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      cout << "#dofs: " << size << endl;
+   }
+
+   ElasticityOperator<finite_stress_qf> hooke(pmesh,
+                                              fespaces->GetFinestFESpace(), matfree,
+                                              dump_matrices);
+
+   if (pmesh.bdr_attributes.Size())
+   {
+      Array<int> ess_attr(pmesh.bdr_attributes.Max());
+      if (problem_type == 0)
+      {
+         ess_attr = 0;
+      }
+      else if (problem_type == 1)
+      {
+         ess_attr = 1;
+      }
+      else if (problem_type == 2)
+      {
+         ess_attr = 0;
+         ess_attr[0] = 1;
+      }
+      else if (problem_type == 3)
+      {
+         ess_attr = 1;
+      }
+      hooke.SetEssentialAttributes(ess_attr);
+      int fixed_dofs_local = hooke.ess_tdof_list.Size();
+      int fixed_dofs_global = 0;
+      MPI_Allreduce(&fixed_dofs_local, &fixed_dofs_global, 1, MPI_INT, MPI_SUM,
+                    pmesh.GetComm());
+      if (Mpi::Root())
+      {
+         out << "#fixed dofs: " << fixed_dofs_global << "\n";
+      }
+   }
+
+   if (problem_type == 1)
+   {
+      // auto body_force = hooke.GetBodyForce();
+      // VectorFunctionCoefficient coeff(2, [](const Vector &coords, Vector &u)
+      // {
+      //    const double x = coords(0);
+      //    const double y = coords(1);
+
+      //    const double a = 0.01, b = 0.05;
+
+      //    const double nu = 0.3;
+      //    const double E = 1.0;
+      //    const double lambda = nu * E / ((1.0 + nu) * (1.0 - 2.0*nu));
+      //    const double mu = E / (2.0 * (1.0 + nu));
+
+      //    u(0) = b * (2.0 * lambda + 2.0 * mu);
+      //    u(1) = a * (2.0 * lambda + 2.0 * mu);
+      // });
+      // body_force->ProjectCoefficient(coeff);
+   }
+
+   ParGridFunction U_gf(&hooke.h1_fes), Ucmp_gf(&hooke.h1_fes);
+   U_gf = 0.0;
+
+   auto boundary_load_ramp = [&](double ramp_scale = 1.0)
+   {
+      Array<int> boundary_load_attr(pmesh.bdr_attributes.Max());
+      boundary_load_attr = 0;
+      boundary_load_attr[1] = 1;
+
+      auto boundary_load = hooke.GetExternalLoad();
+      VectorFunctionCoefficient boundary_load_coeff(2, [&](const Vector &, Vector &u)
+      {
+         u(0) = 0.0;
+         u(1) = -1.0e-3 * ramp_scale;
+      });
+      boundary_load->ProjectBdrCoefficient(boundary_load_coeff, boundary_load_attr);
+   };
+
+   AffineSolution<2> affine_solution;
+   auto patch_test_boundary_load_ramp = [&](ParGridFunction &gf,
+                                            double ramp_scale = 1.0)
+   {
+      u_ex_coeff =
+         std::make_shared<VectorFunctionCoefficient>(2,[&](const Vector &coords,
+                                                           Vector &u)
+      {
+         affine_solution(coords, u);
+         u *= ramp_scale;
+      });
+
+      Array<int> mms_bdr(pmesh.bdr_attributes.Max());
+      mms_bdr = 1;
+      gf.ProjectBdrCoefficient(*u_ex_coeff, mms_bdr);
+   };
+
+   if (problem_type == 1)
+   {
+      u_ex_coeff =
+         std::make_shared<VectorFunctionCoefficient>(2,[](const Vector &coords,
+                                                          Vector &u)
+      {
+         const double x = coords(0);
+         const double y = coords(1);
+         const double a = 0.01, b = 0.05;
+
+         u(0) = a * (2.0 * x + y);
+         u(1) = b * (x + 2.0 * y);
+      });
+
+      Array<int> mms_bdr(pmesh.bdr_attributes.Max());
+      mms_bdr = 1;
+      U_gf.ProjectBdrCoefficient(*u_ex_coeff, mms_bdr);
+   }
+
+   Vector U;
+   U_gf.GetTrueDofs(U);
+
+   if (problem_type == 0)
+   {
+      VectorFunctionCoefficient ucoeff(dim, [&](const Vector &c, Vector &u)
+      {
+         const double x = c(0), y = c(1);
+         u(0) = x*x;
+         u(1) = y;
+         u *= 0.01;
+      });
+      U_gf.ProjectCoefficient(ucoeff);
+      U_gf.GetTrueDofs(U);
+      out << "u: ";
+      U.Print(out, U.Size());
+      Vector R(U.Size());
+      out << "r(u): ";
+      hooke.Mult(U, R);
+      // hooke.Mult(U, R);
+      R.Print(out, U.Size());
+      hooke.GetGradient(U);
+      exit(0);
+   }
+
+   HypreBoomerAMG* amg = nullptr;
+   Multigrid pmg_solver;
+
+   GMRESSolver gmres(MPI_COMM_WORLD);
+   gmres.SetRelTol(1e-8);
+   gmres.SetMaxIter(1000);
+   gmres.SetPrintLevel(2);
+   if (pmg)
+   {
+   }
+   else
+   {
+      if (!matfree)
+      {
+         amg = new HypreBoomerAMG;
+         amg->SetPrintLevel(0);
+         amg->SetElasticityOptions(&hooke.h1_fes);
+         gmres.SetPreconditioner(*amg);
+      }
+   }
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.iterative_mode = true;
+   newton.SetSolver(gmres);
+   newton.SetOperator(hooke);
+   newton.SetRelTol(1e-10);
+   newton.SetMaxIter(50);
+   newton.SetPrintLevel(1);
+
+   if (problem_type == 2)
+   {
+      boundary_load_ramp(1.0);
+
+      Vector zero;
+      newton.Mult(zero, U);
+   }
+   else if (problem_type == 3)
+   {
+      patch_test_boundary_load_ramp(U_gf, 1.0);
+      U_gf.GetTrueDofs(U);
+
+      Vector U_tmp(U), f_tmp(U);
+      auto J = &hooke.GetGradientNoBC(U_tmp);
+
+      J->Mult(U_tmp, f_tmp);
+      f_tmp *= -1.0;
+
+      ParGridFunction f_gf(U_gf);
+      f_gf.Distribute(f_tmp);
+      patch_test_boundary_load_ramp(f_gf, 1.0);
+      f_gf.GetTrueDofs(f_tmp);
+
+      J = &hooke.GetGradient(U_tmp);
+
+      gmres.SetOperator(*J);
+      gmres.Mult(f_tmp, U);
+
+      Vector zero;
+      newton.Mult(zero, U);
+   }
+   else
+   {
+      Vector zero;
+      newton.Mult(zero, U);
+   }
+
+   U_gf.Distribute(U);
+
+   if (problem_type == 1 || problem_type == 3)
+   {
+      out << "||u - u_ex||_L2 = " << U_gf.ComputeL2Error(*u_ex_coeff) << "\n";
+   }
+
+   CALI_MARK_FUNCTION_END;
+   caliper_mgr.flush();
+
+   if (visualization)
+   {
+      if (Mpi::Root())
+      {
+         cout << "writing viz files...\n";
+      }
+      ParaViewDataCollection paraview_dc("hooke", &pmesh);
+      paraview_dc.SetPrefixPath("output");
+      paraview_dc.SetLevelsOfDetail(polynomial_degree);
+      paraview_dc.SetCycle(0);
+      paraview_dc.SetDataFormat(VTKFormat::BINARY);
+      paraview_dc.SetHighOrderOutput(true);
+      paraview_dc.SetTime(0.0); // set the time
+      paraview_dc.RegisterField("displacement", &U_gf);
+      if (problem_type == 1)
+      {
+         Ucmp_gf.ProjectCoefficient(*u_ex_coeff);
+         Ucmp_gf -= U_gf;
+         paraview_dc.RegisterField("displacement_cmp", &Ucmp_gf);
+      }
+      paraview_dc.Save();
+   }
+
+   return 0;
+}

examples/patch2D_quads.mesh

@@ -0,0 +1,42 @@
+MFEM mesh v1.0
+
+#
+# MFEM Geometry Types (see mesh/geom.hpp):
+#
+# POINT       = 0
+# SEGMENT     = 1
+# TRIANGLE    = 2
+# SQUARE      = 3
+# TETRAHEDRON = 4
+# CUBE        = 5
+#
+
+dimension
+2
+
+elements
+5
+1 3 0 1 5 4
+1 3 1 2 6 5
+1 3 2 3 7 6
+1 3 3 0 4 7
+1 3 4 5 6 7
+
+boundary
+4
+1 1 3 0
+2 1 0 1
+3 1 1 2
+4 1 2 3
+
+vertices
+8
+2
+0 0
+1 0
+1 1
+0 1
+0.25 0.3333333333333333
+0.6 0.25
+0.75 0.69
+0.3333333333333333 0.75