renaming and alternative interface

CMakeLists.txt

@@ -14,7 +14,7 @@ set(USER_CONFIG "${CMAKE_CURRENT_SOURCE_DIR}/config/user.cmake" CACHE PATH
   "Path to optional user configuration file.")
 
 # Require C++11 and disable compiler-specific extensions
-set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)
 

miniapps/CMakeLists.txt

@@ -26,3 +26,4 @@ add_subdirectory(tools)
 add_subdirectory(toys)
 add_subdirectory(nurbs)
 add_subdirectory(gslib)
+add_subdirectory(variationalform)

miniapps/variationalform/CMakeLists.txt

@@ -0,0 +1,15 @@
+add_library(parvariationalform parvariationalform.cpp)
+target_include_directories(parvariationalform PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
+target_link_libraries(parvariationalform mfem)
+
+add_executable(parvariationalform_example parvariationalform_example.cpp)
+target_include_directories(parvariationalform_example PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
+target_link_libraries(parvariationalform_example parvariationalform mfem)
+
+add_executable(parvariationalform_ex1 parvariationalform_ex1.cpp)
+target_include_directories(parvariationalform_ex1 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR})
+target_link_libraries(parvariationalform_ex1 parvariationalform mfem)
+
+add_executable(parvariationalform_unit_tests test_tensor_ad.cpp)
+target_include_directories(parvariationalform_unit_tests PRIVATE ${CMAKE_CURRENT_SOURCE_DIR} ${PROJECT_SOURCE_DIR}/tests/unit)
+target_link_libraries(parvariationalform_unit_tests parvariationalform mfem)

miniapps/variationalform/genericintegrator.hpp

@@ -0,0 +1,23 @@
+#include "mfem.hpp"
+
+#pragma once
+
+namespace mfem
+{
+class GenericIntegrator
+{
+public:
+   GenericIntegrator(const IntegrationRule *ir = nullptr) : IntRule(ir) {}
+
+   virtual void Setup(const FiniteElementSpace &) = 0;
+
+   virtual void Apply(const Vector &x, Vector &y) const = 0;
+
+   virtual void ApplyGradient(const Vector &x,
+                              const Vector &v,
+                              Vector &y) const = 0;
+
+protected:
+   const IntegrationRule *IntRule;
+};
+} // namespace mfem

miniapps/variationalform/parvariationalform.cpp

@@ -0,0 +1,94 @@
+#include "parvariationalform.hpp"
+#include "qfuncintegrator.hpp"
+
+namespace mfem
+{
+ParVariationalForm::ParVariationalForm(ParFiniteElementSpace *f)
+   : Operator(f->GetTrueVSize()), fes(f), P(f->GetProlongationMatrix()),
+     grad(*this)
+{
+   G = fes->GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC);
+   MFEM_ASSERT(G, "Some GetElementRestriction error");
+   x_local.SetSize(G->Height(), Device::GetMemoryType());
+   v_local.SetSize(G->Height(), Device::GetMemoryType());
+   y_local.SetSize(G->Height(), Device::GetMemoryType());
+}
+
+void ParVariationalForm::Mult(const Vector &x, Vector &y) const
+{
+   px.SetSize(P->Height());
+   py.SetSize(P->Height());
+
+   P->Mult(x, px);
+   G->Mult(px, x_local);
+
+   y_local = 0.0;
+   for (int i = 0; i < domain_integrators.Size(); ++i)
+   {
+      // y += F(x)
+      domain_integrators[i]->Apply(x_local, y_local);
+   }
+
+   G->MultTranspose(y_local, py);
+   P->MultTranspose(py, y);
+
+   y.HostReadWrite();
+   for (int i = 0; i < ess_tdof_list.Size(); i++)
+   {
+      y(ess_tdof_list[i]) = 0.0;
+   }
+}
+
+void ParVariationalForm::GradientMult(const Vector &v, Vector &y) const
+{
+   px.SetSize(P->Height());
+   py.SetSize(P->Height());
+   pv.SetSize(P->Height());
+
+   P->Mult(v, pv);
+   G->Mult(pv, v_local);
+
+   P->Mult(x_lin, px);
+   G->Mult(px, x_local);
+
+   y_local = 0.0;
+   for (int i = 0; i < domain_integrators.Size(); ++i)
+   {
+      // y += dF(x)/dx * v
+      if (is_linear)
+      {
+         // take care of RHS
+         // domain_integrators[i]->Apply(v_local, y_local);
+      }
+      domain_integrators[i]->ApplyGradient(x_local, v_local, y_local);
+   }
+
+   G->MultTranspose(y_local, py);
+   P->MultTranspose(py, y);
+
+   y.HostReadWrite();
+   for (int i = 0; i < ess_tdof_list.Size(); i++)
+   {
+      y(ess_tdof_list[i]) = v(ess_tdof_list[i]);
+   }
+}
+
+Operator &ParVariationalForm::GetGradient(const Vector &x) const
+{
+   x_lin = x;
+   return grad;
+}
+
+HypreParMatrix *ParVariationalForm::GetGradientMatrix(const Vector &x)
+{
+   delete gradient_matrix;
+   gradient_matrix = new HypreParMatrix;
+   return gradient_matrix;
+}
+
+void ParVariationalForm::SetEssentialBC(const Array<int> &ess_attr)
+{
+   fes->GetEssentialTrueDofs(ess_attr, ess_tdof_list);
+}
+
+} // namespace mfem

miniapps/variationalform/parvariationalform.hpp

@@ -0,0 +1,109 @@
+#include "mfem.hpp"
+#include "genericintegrator.hpp"
+#include "qfuncintegrator.hpp"
+
+#pragma once
+
+namespace mfem
+{
+class ParVariationalForm : public Operator
+{
+   class Gradient : public Operator
+   {
+   public:
+      Gradient(ParVariationalForm &f) : Operator(f.Height()), form(f){};
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         form.GradientMult(x, y);
+      }
+
+   private:
+      ParVariationalForm &form;
+   };
+
+public:
+   ParVariationalForm(ParFiniteElementSpace *f);
+
+   void Mult(const Vector &x, Vector &y) const override;
+
+   void AddDomainIntegrator(GenericIntegrator *i)
+   {
+      domain_integrators.Append(i);
+      i->Setup(*fes);
+   }
+
+   template<typename integrator_type,
+            typename qfunc_type,
+            typename... qfunc_args_type>
+   void AddDomainIntegrator(qfunc_type f, qfunc_args_type const &... fargs)
+   {
+      if constexpr (std::is_same_v<integrator_type, DomainLFIntegrator>)
+      {
+         auto i = new QFunctionIntegrator(
+            [&](auto... args) {
+               auto du = std::get<1>(std::tuple{args...});
+               return std::tuple{f(args...), decltype(du){}};
+            },
+            0,
+            *fes->GetParMesh(),
+            fargs...);
+
+         domain_integrators.Append(i);
+         i->Setup(*fes);
+      }
+
+      if constexpr (std::is_same_v<integrator_type, DiffusionIntegrator>)
+      {
+         auto i = new QFunctionIntegrator(
+            [&](auto... args) {
+               auto du = std::get<1>(std::tuple{args...});
+               return std::tuple{0.0, f(args...) * du};
+            },
+            0,
+            *fes->GetParMesh(),
+            fargs...);
+
+         domain_integrators.Append(i);
+         i->Setup(*fes);
+      }
+   }
+
+   void AssumeLinear() { is_linear = true; }
+
+   // Return an Operator that provides a Mult(x, y) which is the MatVec of the
+   // gradient of the ParVariationalForm wrt x. Acts as a "passthrough" to
+   // ::GradientMult in order to satisfy mfem interfaces.
+   Operator &GetGradient(const Vector &x) const override;
+
+   // Return an assmbled parallel matrix which represents the gradient of the
+   // ParVariationalForm wrt to x.
+   HypreParMatrix *GetGradientMatrix(const Vector &x);
+
+   void SetEssentialBC(const Array<int> &ess_attr);
+
+protected:
+   // y = F'(x_lin) * v
+   void GradientMult(const Vector &v, Vector &y) const;
+
+   bool is_linear = false;
+   ParFiniteElementSpace *fes;
+   Array<GenericIntegrator *> domain_integrators;
+   Array<int> ess_tdof_list;
+
+   // T -> L
+   const Operator *P;
+
+   // L -> E
+   const Operator *G;
+
+   mutable Vector x_local, y_local, v_local, px, py, pv;
+
+   // State to build the Gradient on, single source of the true state.
+   mutable Vector x_lin;
+
+   mutable Gradient grad;
+
+   HypreParMatrix *gradient_matrix = nullptr;
+};
+} // namespace mfem

miniapps/variationalform/parvariationalform_ex1.cpp

@@ -0,0 +1,164 @@
+#include "mfem.hpp"
+#include "parvariationalform.hpp"
+#include "tensor.hpp"
+#include <fstream>
+#include <iostream>
+
+using namespace std;
+using namespace mfem;
+
+int main(int argc, char *argv[])
+{
+   int num_procs, myid;
+   MPI_Init(&argc, &argv);
+   MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
+   MPI_Comm_rank(MPI_COMM_WORLD, &myid);
+
+   const char *mesh_file = "../data/inline-quad.mesh";
+   int order = 1;
+   bool static_cond = false;
+   bool pa = false;
+   const char *device_config = "cpu";
+   bool visualization = true;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&order,
+                  "-o",
+                  "--order",
+                  "Finite element order (polynomial degree) or -1 for"
+                  " isoparametric space.");
+   args.AddOption(&visualization,
+                  "-vis",
+                  "--visualization",
+                  "-no-vis",
+                  "--no-visualization",
+                  "Enable or disable GLVis visualization.");
+   args.Parse();
+   if (!args.Good())
+   {
+      if (myid == 0)
+      {
+         args.PrintUsage(cout);
+      }
+      MPI_Finalize();
+      return 1;
+   }
+   if (myid == 0)
+   {
+      args.PrintOptions(cout);
+   }
+
+   Mesh mesh(mesh_file, 1, 1);
+   int dim = mesh.Dimension();
+
+   {
+      int ref_levels = (int) floor(log(10. / mesh.GetNE()) / log(2.) / dim);
+      for (int l = 0; l < ref_levels; l++)
+      {
+         mesh.UniformRefinement();
+      }
+   }
+
+   ParMesh pmesh(MPI_COMM_WORLD, mesh);
+   mesh.Clear();
+   {
+      int par_ref_levels = 2;
+      for (int l = 0; l < par_ref_levels; l++)
+      {
+         pmesh.UniformRefinement();
+      }
+   }
+
+   FiniteElementCollection *fec;
+   bool delete_fec;
+   if (order > 0)
+   {
+      fec = new H1_FECollection(order, dim);
+      delete_fec = true;
+   }
+   else if (pmesh.GetNodes())
+   {
+      fec = pmesh.GetNodes()->OwnFEC();
+      delete_fec = false;
+      if (myid == 0)
+      {
+         cout << "Using isoparametric FEs: " << fec->Name() << endl;
+      }
+   }
+   else
+   {
+      fec = new H1_FECollection(order = 1, dim);
+      delete_fec = true;
+   }
+   ParFiniteElementSpace fespace(&pmesh, fec);
+   HYPRE_Int size = fespace.GlobalTrueVSize();
+   if (myid == 0)
+   {
+      cout << "Number of finite element unknowns: " << size << endl;
+   }
+
+   Array<int> ess_bdr(pmesh.bdr_attributes.Max());
+   ess_bdr = 1;
+
+   FunctionCoefficient u_excoeff([&](const Vector &coords) {
+      double x = coords(0);
+      double y = coords(1);
+      return x * x + y * y;
+   });
+
+   ParGridFunction x(&fespace);
+
+   x.ProjectBdrCoefficient(u_excoeff, ess_bdr);
+
+   ParVariationalForm form(&fespace);
+
+   auto b_coeff = [&](auto u, auto du, auto x) {
+      return 4.0 * (1.0 + 2.0 * x[0] * x[0] + 2.0 * x[1] * x[1]);
+   };
+
+   form.AddDomainIntegrator<DomainLFIntegrator>(b_coeff);
+
+   auto a_coeff = [&](auto u, auto du, auto x) { return 1.0 + u; };
+
+   form.AddDomainIntegrator<DiffusionIntegrator>(a_coeff);
+
+   form.SetEssentialBC(ess_bdr);
+
+   CGSolver cg(MPI_COMM_WORLD);
+   cg.SetRelTol(1e-6);
+   cg.SetMaxIter(2000);
+   cg.SetPrintLevel(1);
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.SetOperator(form);
+   newton.SetSolver(cg);
+   newton.SetPrintLevel(1);
+   newton.SetRelTol(1e-8);
+   newton.SetMaxIter(100);
+
+   Vector zero;
+   Vector X;
+   x.GetTrueDofs(X);
+   newton.Mult(zero, X);
+
+   x.Distribute(X);
+
+   if (visualization)
+   {
+      char vishost[] = "localhost";
+      int visport = 19916;
+      socketstream sol_sock(vishost, visport);
+      sol_sock << "parallel " << num_procs << " " << myid << "\n";
+      sol_sock.precision(8);
+      sol_sock << "solution\n" << pmesh << x << flush;
+   }
+
+   if (delete_fec)
+   {
+      delete fec;
+   }
+   MPI_Finalize();
+
+   return 0;
+}

miniapps/variationalform/parvariationalform_example.cpp

@@ -0,0 +1,118 @@
+#include "mfem.hpp"
+#include "parvariationalform.hpp"
+#include "qfuncintegrator.hpp"
+#include "tensor.hpp"
+#include <fstream>
+#include <iostream>
+
+using namespace std;
+using namespace mfem;
+
+int main(int argc, char *argv[])
+{
+   int num_procs, myid;
+   MPI_Init(&argc, &argv);
+   MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
+   MPI_Comm_rank(MPI_COMM_WORLD, &myid);
+
+   const char *mesh_file = "../data/inline-quad.mesh";
+   int order = 1;
+   int refinements = 0;
+   double p = 5.0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&refinements, "-r", "--ref", "");
+   args.AddOption(&order, "-o", "--order", "");
+
+   args.Parse();
+   if (!args.Good())
+   {
+      if (myid == 0)
+      {
+         args.PrintUsage(cout);
+      }
+      MPI_Finalize();
+      return 1;
+   }
+   if (myid == 0)
+   {
+      args.PrintOptions(cout);
+   }
+
+   Mesh mesh(mesh_file, 1, 1);
+
+   int dim = mesh.Dimension();
+   {
+      for (int l = 0; l < refinements; l++)
+      {
+         mesh.UniformRefinement();
+      }
+   }
+
+   ParMesh pmesh(MPI_COMM_WORLD, mesh);
+
+   auto fec = H1_FECollection(order, dim);
+   ParFiniteElementSpace fespace(&pmesh, &fec);
+
+   Array<int> ess_bdr(pmesh.bdr_attributes.Max());
+   ess_bdr = 1;
+
+   FunctionCoefficient u_excoeff([&](const Vector &coords) {
+      double x = coords(0);
+      double y = coords(1);
+
+      // return 1.0 - pow(sqrt(x * x + y * y), p / (p - 1.0));
+      return x * x + y * y;
+   });
+
+   ParGridFunction x(&fespace);
+   x.Randomize();
+
+   x.ProjectBdrCoefficient(u_excoeff, ess_bdr);
+
+   ParVariationalForm form(&fespace);
+
+   auto plaplacian = new QFunctionIntegrator([&](auto u, auto du, auto x) {
+      // auto f0 = -1.0;
+      // auto f1 = pow(norm(du), p - 2.0) * du;
+      auto f0 = 4.0 * (1.0 + 2.0 * x[0] * x[0] + 2.0 * x[1] * x[1]);
+      auto f1 = (1.0 + u) * du;
+      return std::tuple{f0, f1};
+   }, 0, pmesh);
+
+   form.AddDomainIntegrator(plaplacian);
+
+   form.SetEssentialBC(ess_bdr);
+
+   CGSolver cg(MPI_COMM_WORLD);
+   cg.SetRelTol(1e-6);
+   cg.SetMaxIter(2000);
+   cg.SetPrintLevel(1);
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.SetOperator(form);
+   newton.SetSolver(cg);
+   newton.SetPrintLevel(1);
+   newton.SetRelTol(1e-8);
+   newton.SetMaxIter(100);
+
+   Vector zero;
+   Vector X;
+   x.GetTrueDofs(X);
+   newton.Mult(zero, X);
+
+   x.Distribute(X);
+
+   // x.ProjectCoefficient(u_excoeff);
+   
+   char vishost[] = "localhost";
+   int visport = 19916;
+   socketstream sol_sock(vishost, visport);
+   sol_sock << "parallel " << num_procs << " " << myid << "\n";
+   sol_sock.precision(8);
+   sol_sock << "solution\n" << pmesh << x << flush;
+
+   MPI_Finalize();
+
+   return 0;
+}

miniapps/variationalform/qfuncintegrator.hpp

@@ -0,0 +1,404 @@
+#include "mfem.hpp"
+#include "../../general/forall.hpp"
+#include "genericintegrator.hpp"
+#include "tensor.hpp"
+
+#pragma once
+
+namespace mfem
+{
+template<typename T>
+struct supported_type
+{
+   static constexpr bool value = false;
+};
+
+template<>
+struct supported_type<ParMesh>
+{
+   static constexpr bool value = true;
+};
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+class QFunctionIntegrator : public GenericIntegrator
+{
+public:
+   QFunctionIntegrator(qfunc_type f,
+                       qfunc_grad_type f_grad,
+                       qfunc_args_type const &... fargs);
+
+   QFunctionIntegrator(qfunc_type f, qfunc_args_type const &... fargs);
+
+   void Setup(const FiniteElementSpace &fes) override;
+
+   void Apply(const Vector &, Vector &) const override;
+
+   // y += F'(x) * v
+   void ApplyGradient(const Vector &x,
+                      const Vector &v,
+                      Vector &y) const override;
+
+protected:
+   template<int D1D, int Q1D>
+   void Apply2D(const Vector &u_in_, Vector &y_) const;
+
+   template<int D1D, int Q1D>
+   void ApplyGradient2D(const Vector &u_in_,
+                        const Vector &v_in_,
+                        Vector &y_) const;
+
+   auto EvaluateFargValue(const Mesh &m, const double qx, const double qy) const;
+
+   const FiniteElementSpace *fespace;
+   const DofToQuad *maps;        ///< Not owned
+   const GeometricFactors *geom; ///< Not owned
+   int dim, ne, nq, dofs1D, quad1D;
+
+   // Geometric factors
+   Vector J_;
+   Vector W_;
+
+   qfunc_type qf;
+   qfunc_grad_type qf_grad;
+   std::tuple<qfunc_args_type...> qf_farg_values;
+};
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::
+   QFunctionIntegrator(qfunc_type f,
+                       qfunc_grad_type df,
+                       qfunc_args_type const &... fargs)
+   : GenericIntegrator(nullptr), maps(nullptr), geom(nullptr), qf(f),
+     qf_grad(df), qf_farg_values(std::tuple{fargs...})
+{
+   static_assert((supported_type<qfunc_args_type>::value && ...),
+                 "Type not supported for parameter expansion. See "
+                 "documentation for supported types.");
+}
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::
+   QFunctionIntegrator(qfunc_type f, qfunc_args_type const &... fargs)
+   : GenericIntegrator(nullptr), maps(nullptr), geom(nullptr), qf(f),
+     qf_grad(qfunc_grad_type{}), qf_farg_values(std::tuple{fargs...})
+{
+   static_assert((supported_type<qfunc_args_type>::value && ...),
+                 "Type not supported for parameter expansion. See "
+                 "documentation for supported types.");
+}
+
+template<typename qfunc_type, typename... qfunc_args_type>
+QFunctionIntegrator(qfunc_type, qfunc_args_type const &...)
+   -> QFunctionIntegrator<qfunc_type, int, qfunc_args_type const &...>;
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+void QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::Setup(
+   const FiniteElementSpace &fes)
+{
+   // Assuming the same element type
+   fespace = &fes;
+   Mesh *mesh = fes.GetMesh();
+   if (mesh->GetNE() == 0)
+   {
+      return;
+   }
+   const FiniteElement &el = *fes.GetFE(0);
+   ElementTransformation *T = mesh->GetElementTransformation(0);
+   const IntegrationRule *ir = nullptr;
+   if (!IntRule)
+   {
+      IntRule = &IntRules.Get(el.GetGeomType(), el.GetOrder() * 2);
+   }
+   ir = IntRule;
+
+   dim = mesh->Dimension();
+   ne = fes.GetMesh()->GetNE();
+   nq = ir->GetNPoints();
+   geom = mesh->GetGeometricFactors(*ir,
+                                    GeometricFactors::COORDINATES
+                                       | GeometricFactors::JACOBIANS);
+   maps = &el.GetDofToQuad(*ir, DofToQuad::TENSOR);
+   dofs1D = maps->ndof;
+   quad1D = maps->nqpt;
+   //    pa_data.SetSize(ne * nq, Device::GetDeviceMemoryType());
+
+   W_.SetSize(nq, Device::GetDeviceMemoryType());
+   W_.GetMemory().CopyFrom(ir->GetWeights().GetMemory(), nq);
+
+   // J.SetSize(ne * nq, Device::GetDeviceMemoryType());
+   J_ = geom->J;
+}
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+auto QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::
+   EvaluateFargValue(const Mesh &m, const double qx, const double qy) const
+{
+   Vector trip(3);
+   trip = 0.0;
+   ElementTransformation *tr = const_cast<Mesh &>(m).GetElementTransformation(0);
+   tr->Transform(IntRule->IntPoint(qx + quad1D * qy), trip);
+   return tensor<double, 3>{
+      {trip(0), trip(1), m.SpaceDimension() == 2 ? 0.0 : trip(2)}};
+}
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+void QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::Apply(
+   const Vector &x, Vector &y) const
+{
+   if (dim == 2)
+   {
+      switch ((dofs1D << 4) | quad1D)
+      {
+      case 0x22:
+         return Apply2D<2, 2>(x, y);
+      default:
+         MFEM_ASSERT(false, "NOPE");
+      }
+   }
+}
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+template<int D1D, int Q1D>
+void QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::Apply2D(
+   const Vector &u_in_, Vector &y_) const
+{
+   int NE = ne;
+
+   auto v1d = Reshape(maps->B.Read(), Q1D, D1D);
+   auto dv1d_dX = Reshape(maps->G.Read(), Q1D, D1D);
+   // (NQ x SDIM x DIM x NE)
+   auto J = Reshape(J_.Read(), Q1D, Q1D, 2, 2, NE);
+   auto W = Reshape(W_.Read(), Q1D, Q1D);
+   auto u = Reshape(u_in_.Read(), D1D, D1D, NE);
+   auto y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
+
+   // MFEM_FORALL(e, NE, {
+   for (int e = 0; e < NE; e++)
+   {
+      // loop over quadrature points
+      for (int qy = 0; qy < Q1D; ++qy)
+      {
+         for (int qx = 0; qx < Q1D; ++qx)
+         {
+            double u_q = 0.0;
+            double du_dX_q[2] = {0.0};
+            for (int ix = 0; ix < D1D; ix++)
+            {
+               for (int iy = 0; iy < D1D; iy++)
+               {
+                  u_q += u(ix, iy, e) * v1d(qx, ix) * v1d(qy, iy);
+                  du_dX_q[0] += u(ix, iy, e) * dv1d_dX(qx, ix) * v1d(qy, iy);
+                  du_dX_q[1] += u(ix, iy, e) * v1d(qx, ix) * dv1d_dX(qy, iy);
+               }
+            }
+
+            // du_dx_q = invJ^T * du_dX_q
+            //         = (adjJ^T * du_dX_q) / detJ
+            double J_q[2][2] = {{J(qx, qy, 0, 0, e),
+                                 J(qx, qy, 0, 1, e)}, // J_q[0][0], J_q[0][1]
+                                {J(qx, qy, 1, 0, e),
+                                 J(qx, qy, 1, 1, e)}}; // J_q[1][0], J_q[1][1]
+
+            double detJ_q = (J_q[0][0] * J_q[1][1]) - (J_q[0][1] * J_q[1][0]);
+
+            double adjJ[2][2] = {{J_q[1][1], -J_q[0][1]},
+                                 {-J_q[1][0], J_q[0][0]}};
+
+            tensor<double, 2> du_dx_q
+               = {(adjJ[0][0] * du_dX_q[0] + adjJ[1][0] * du_dX_q[1]) / detJ_q,
+                  (adjJ[0][1] * du_dX_q[0] + adjJ[1][1] * du_dX_q[1]) / detJ_q};
+
+            auto processed_qf_farg_values = std::apply(
+               [=](auto &... a) {
+                  return std::make_tuple(u_q,
+                                         du_dx_q,
+                                         EvaluateFargValue(a, qx, qy)...);
+               },
+               qf_farg_values);
+
+            auto [f0, f1] = std::apply(qf, processed_qf_farg_values);
+
+            double f0_X = f0 * detJ_q;
+
+            // f1_X = invJ * f1 * detJ
+            //      = adjJ * f1
+            double f1_X[2] = {
+               adjJ[0][0] * f1[0] + adjJ[0][1] * f1[1],
+               adjJ[1][0] * f1[0] + adjJ[1][1] * f1[1],
+            };
+
+            for (int ix = 0; ix < D1D; ix++)
+            {
+               for (int iy = 0; iy < D1D; iy++)
+               {
+                  // accumulate v * f0 + dot(dv_dx, f1)
+                  y(ix, iy, e) += (f0_X * v1d(qx, ix) * v1d(qy, iy)
+                                   + f1_X[0] * dv1d_dX(qx, ix) * v1d(qy, iy)
+                                   + f1_X[1] * dv1d_dX(qy, iy) * v1d(qx, ix))
+                                  * W(qx, qy);
+               }
+            }
+         }
+      }
+   }
+   // });
+}
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+void QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::
+   ApplyGradient(const Vector &x, const Vector &v, Vector &y) const
+{
+   ApplyGradient2D<2, 2>(x, v, y);
+}
+
+template<typename qfunc_type, typename qfunc_grad_type, typename... qfunc_args_type>
+template<int D1D, int Q1D>
+void QFunctionIntegrator<qfunc_type, qfunc_grad_type, qfunc_args_type...>::
+   ApplyGradient2D(const Vector &u_in_, const Vector &v_in_, Vector &y_) const
+{
+   int NE = ne;
+
+   auto v1d = Reshape(maps->B.Read(), Q1D, D1D);
+   auto dv1d_dX = Reshape(maps->G.Read(), Q1D, D1D);
+   // (NQ x SDIM x DIM x NE)
+   auto J = Reshape(J_.Read(), Q1D, Q1D, 2, 2, NE);
+   auto W = Reshape(W_.Read(), Q1D, Q1D);
+   auto u = Reshape(u_in_.Read(), D1D, D1D, NE);
+   auto v = Reshape(v_in_.Read(), D1D, D1D, NE);
+   auto y = Reshape(y_.ReadWrite(), D1D, D1D, NE);
+
+   for (int e = 0; e < NE; e++)
+   {
+      // loop over quadrature points
+      for (int qy = 0; qy < Q1D; ++qy)
+      {
+         for (int qx = 0; qx < Q1D; ++qx)
+         {
+            double u_q = 0.0;
+            double du_dX_q[2] = {0.0};
+            double v_q = 0.0;
+            double dv_dX_q[2] = {0.0};
+            for (int ix = 0; ix < D1D; ix++)
+            {
+               for (int iy = 0; iy < D1D; iy++)
+               {
+                  u_q += u(ix, iy, e) * v1d(qx, ix) * v1d(qy, iy);
+
+                  du_dX_q[0] += u(ix, iy, e) * dv1d_dX(qx, ix) * v1d(qy, iy);
+                  du_dX_q[1] += u(ix, iy, e) * v1d(qx, ix) * dv1d_dX(qy, iy);
+
+                  v_q += v(ix, iy, e) * v1d(qx, ix) * v1d(qy, iy);
+
+                  dv_dX_q[0] += v(ix, iy, e) * dv1d_dX(qx, ix) * v1d(qy, iy);
+                  dv_dX_q[1] += v(ix, iy, e) * v1d(qx, ix) * dv1d_dX(qy, iy);
+               }
+            }
+
+            // du_dx_q = invJ^T * du_dX_q
+            //         = (adjJ^T * du_dX_q) / detJ
+            double J_q[2][2] = {{J(qx, qy, 0, 0, e),
+                                 J(qx, qy, 0, 1, e)}, // J_q[0][0], J_q[0][1]
+                                {J(qx, qy, 1, 0, e),
+                                 J(qx, qy, 1, 1, e)}}; // J_q[1][0], J_q[1][1]
+
+            double detJ_q = (J_q[0][0] * J_q[1][1]) - (J_q[0][1] * J_q[1][0]);
+
+            double adjJ[2][2] = {{J_q[1][1], -J_q[0][1]},
+                                 {-J_q[1][0], J_q[0][0]}};
+
+            tensor<double, 2> du_dx_q
+               = {(adjJ[0][0] * du_dX_q[0] + adjJ[1][0] * du_dX_q[1]) / detJ_q,
+                  (adjJ[0][1] * du_dX_q[0] + adjJ[1][1] * du_dX_q[1]) / detJ_q};
+
+            double dv_dx_q[2]
+               = {(adjJ[0][0] * dv_dX_q[0] + adjJ[1][0] * dv_dX_q[1]) / detJ_q,
+                  (adjJ[0][1] * dv_dX_q[0] + adjJ[1][1] * dv_dX_q[1]) / detJ_q};
+
+            // compute dF(u, du)/du
+            auto processed_qf_farg_values_u = std::apply(
+               [=](auto &... a) {
+                  return std::make_tuple(derivative_wrt(u_q),
+                                         du_dx_q,
+                                         EvaluateFargValue(a, qx, qy)...);
+               },
+               qf_farg_values);
+
+            auto [f0u, f1u] = std::apply(qf, processed_qf_farg_values_u);
+
+            double f00 = 0.0;
+            if constexpr (std::is_same_v<decltype(f0u), double>)
+            {
+               f00 = 0.0;
+            }
+            else
+            {
+               f00 = f0u.gradient;
+            }
+
+            tensor<double, 2> f10;
+            if constexpr (std::is_same_v<decltype(f1u), tensor<double, 2>>)
+            {
+               f10 = {0.0, 0.0};
+            }
+            else
+            {
+               f10 = {f1u[0].gradient, f1u[1].gradient};
+            }
+
+            // compute dF(u, du)/ddu
+            auto processed_qf_farg_values_du = std::apply(
+               [=](auto &... a) {
+                  return std::make_tuple(u_q,
+                                         derivative_wrt(du_dx_q),
+                                         EvaluateFargValue(a, qx, qy)...);
+               },
+               qf_farg_values);
+
+            auto [f0du, f1du] = std::apply(qf, processed_qf_farg_values_du);
+
+            tensor<double, 2> f01;
+            if constexpr (std::is_same_v<decltype(f0du), double>)
+            {
+               f01 = {0.0, 0.0};
+            }
+            else
+            {
+               f01 = f0du.gradient;
+            }
+            tensor<double, 2, 2> f11 = {
+               {{f1du[0].gradient[0], f1du[0].gradient[1]},
+                {f1du[1].gradient[0], f1du[1].gradient[1]}}};
+
+            double W0 = f00 * v_q + f01[0] * dv_dx_q[0] + f01[1] * dv_dx_q[1];
+
+            double W1[2] = {f10[0] * v_q + f11[0][0] * dv_dx_q[0]
+                               + f11[0][1] * dv_dx_q[1],
+                            f10[1] * v_q + f11[1][0] * dv_dx_q[0]
+                               + +f11[1][1] * dv_dx_q[1]};
+
+            double W0_X = W0 * detJ_q;
+
+            // W1_X = invJ * W1 * detJ
+            //      = adjJ * W1
+            double W1_X[2] = {
+               adjJ[0][0] * W1[0] + adjJ[0][1] * W1[1],
+               adjJ[1][0] * W1[0] + adjJ[1][1] * W1[1],
+            };
+
+            for (int ix = 0; ix < D1D; ix++)
+            {
+               for (int iy = 0; iy < D1D; iy++)
+               {
+                  // @TODO: proper comment
+                  y(ix, iy, e) += (W0_X * v1d(qx, ix) * v1d(qy, iy)
+                                   + W1_X[0] * dv1d_dX(qx, ix) * v1d(qy, iy)
+                                   + W1_X[1] * dv1d_dX(qy, iy) * v1d(qx, ix))
+                                  * W(qx, qy);
+               }
+            }
+         }
+      }
+   }
+}
+
+} // namespace mfem

miniapps/variationalform/tensor.hpp

@@ -0,0 +1,347 @@
+// tensor
+#include <iostream>
+
+#pragma once
+
+template<typename T, int... n>
+struct tensor;
+
+template<typename T>
+struct tensor<T, 1>
+{
+   static constexpr int shape[1] = {1};
+   operator T() { return value; }
+   T value;
+};
+
+template<typename T, int n>
+struct tensor<T, n>
+{
+   static constexpr int shape[1] = {n};
+   constexpr auto &operator[](int i) { return value[i]; };
+   constexpr auto operator[](int i) const { return value[i]; };
+   T value[n];
+};
+
+template<typename T, int first, int... rest>
+struct tensor<T, first, rest...>
+{
+   static constexpr int shape[1 + sizeof...(rest)] = {first, rest...};
+   constexpr auto &operator[](int i) { return value[i]; };
+   constexpr auto operator[](int i) const { return value[i]; };
+   tensor<T, rest...> value[first];
+};
+
+template<int n>
+constexpr int product(int (&values)[n])
+{
+   int p = 1;
+   for (int i = 0; i < n; i++)
+   {
+      p *= values[i];
+   }
+   return p;
+}
+
+template<typename S, typename T, int... n>
+auto operator+(tensor<S, n...> A, tensor<T, n...> B)
+{
+   tensor<decltype(S{} + T{}), n...> C{};
+   for (int i = 0; i < tensor<T, n...>::shape[0]; i++)
+   {
+      C[i] = A[i] + B[i];
+   }
+   return C;
+}
+
+template<typename S, typename T, int... n>
+auto operator-(tensor<S, n...> A, tensor<T, n...> B)
+{
+   tensor<decltype(S{} - T{}), n...> C{};
+   for (int i = 0; i < tensor<T, n...>::shape[0]; i++)
+   {
+      C[i] = A[i] - B[i];
+   }
+   return C;
+}
+
+template<typename S, typename T, int... n>
+auto operator*(S scale, tensor<T, n...> A)
+{
+   tensor<decltype(S{} * T{}), n...> C{};
+   for (int i = 0; i < tensor<T, n...>::shape[0]; i++)
+   {
+      C[i] = scale * A[i];
+   }
+   return C;
+}
+
+template<typename S, typename T, int... n>
+auto operator*(tensor<T, n...> A, S scale)
+{
+   tensor<decltype(T{} * S{}), n...> C{};
+   for (int i = 0; i < tensor<T, n...>::shape[0]; i++)
+   {
+      C[i] = A[i] * scale;
+   }
+   return C;
+}
+
+template<typename S, typename T, int... n>
+auto operator/(S scale, tensor<T, n...> A)
+{
+   tensor<decltype(S{} / T{}), n...> C{};
+   for (int i = 0; i < tensor<T, n...>::shape[0]; i++)
+   {
+      C[i] = scale / A[i];
+   }
+   return C;
+}
+
+template<typename S, typename T, int... n>
+auto operator/(tensor<T, n...> A, S scale)
+{
+   tensor<decltype(T{} / S{}), n...> C{};
+   for (int i = 0; i < tensor<T, n...>::shape[0]; i++)
+   {
+      C[i] = A[i] / scale;
+   }
+   return C;
+}
+
+template<typename S, typename T, int m, int n, int p>
+auto dot(tensor<S, m, n> A, tensor<T, n, p> B)
+{
+   tensor<decltype(S{} * T{}), m, p> AB{};
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < p; j++)
+      {
+         for (int k = 0; k < n; k++)
+         {
+            AB[i][j] = AB[i][j] + A[i][k] * B[k][j];
+         }
+      }
+   }
+   return AB;
+}
+
+template<typename T, int m, int n>
+auto inner(tensor<T, m, n> A, tensor<T, m, n> B)
+{
+   double value = 0.0;
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         value += A[i][j] * B[i][j];
+      }
+   }
+   return value;
+}
+
+auto inner(double a, double b)
+{
+   return a * b;
+}
+
+template<typename T, int n>
+auto tr(tensor<T, n, n> A)
+{
+   T trA{};
+   for (int i = 0; i < n; i++)
+   {
+      trA = trA + A[i][i];
+   }
+   return trA;
+}
+
+template<int dim>
+constexpr tensor<double, dim, dim> Identity()
+{
+   tensor<double, dim, dim> I{};
+   for (int i = 0; i < dim; i++)
+   {
+      for (int j = 0; j < dim; j++)
+      {
+         I[i][j] = (i == j);
+      }
+   }
+   return I;
+}
+
+template<typename T, int m, int n>
+auto transpose(const tensor<T, m, n> &A)
+{
+   tensor<T, n, m> AT{};
+   for (int i = 0; i < n; i++)
+   {
+      for (int j = 0; j < m; j++)
+      {
+         AT[i][j] = A[j][i];
+      }
+   }
+   return AT;
+}
+
+template<typename T, int n>
+auto norm(const tensor<T, n> &A)
+{
+   T r = {};
+   for (int i = 0; i < n; i++)
+   {
+      r = r + A[i] * A[i];
+   }
+   return pow(r, 0.5);
+}
+
+template<typename gradient_type>
+struct dual
+{
+   double value;
+   gradient_type gradient;
+};
+
+template<typename gradient_type>
+auto operator+(dual<gradient_type> a, double b)
+{
+   return dual<gradient_type>{a.value + b, a.gradient};
+}
+
+template<typename gradient_type>
+auto operator+(double a, dual<gradient_type> b)
+{
+   return dual<gradient_type>{a + b.value, b.gradient};
+}
+
+template<typename gradient_type>
+auto operator+(dual<gradient_type> a, dual<gradient_type> b)
+{
+   return dual<gradient_type>{a.value + b.value, a.gradient + b.gradient};
+}
+
+template<typename gradient_type>
+auto operator*(dual<gradient_type> a, double b)
+{
+   return dual<gradient_type>{a.value * b, a.gradient * b};
+}
+
+template<typename gradient_type>
+auto operator*(double a, dual<gradient_type> b)
+{
+   return dual<gradient_type>{a * b.value, a * b.gradient};
+}
+
+template<typename gradient_type>
+auto operator*(dual<gradient_type> a, dual<gradient_type> b)
+{
+   return dual<gradient_type>{a.value * b.value,
+                              b.value * a.gradient + a.value * b.gradient};
+}
+
+template<typename gradient_type>
+auto cos(dual<gradient_type> a)
+{
+   return dual<gradient_type>{cos(a.value), -a.gradient * sin(a.value)};
+}
+
+template<typename gradient_type>
+auto exp(dual<gradient_type> a)
+{
+   return dual<gradient_type>{exp(a.value), exp(a.value)};
+}
+
+template<typename gradient_type>
+auto log(dual<gradient_type> a)
+{
+   return dual<gradient_type>{log(a.value), a.gradient / a.value};
+}
+
+template<typename gradient_type>
+auto pow(dual<gradient_type> a, dual<gradient_type> b)
+{
+   double value = pow(a.value, b.value);
+   return dual<gradient_type>{value,
+                              value
+                                 * (a.gradient * (b.value / a.value)
+                                    + b.gradient * log(a.value))};
+}
+
+template<typename gradient_type>
+auto pow(double a, dual<gradient_type> b)
+{
+   double value = pow(a, b.value);
+   return dual<gradient_type>{value, value * b.gradient * log(a)};
+}
+
+template<typename gradient_type>
+auto pow(dual<gradient_type> a, double b)
+{
+   double value = pow(a.value, b);
+   return dual<gradient_type>{value, value * a.gradient * b / a.value};
+}
+
+template<typename T, int... n>
+auto &operator<<(std::ostream &out, dual<T> A)
+{
+   out << '(' << A.value << ' ' << A.gradient << ')';
+   return out;
+}
+
+template<typename T, int... n>
+auto &operator<<(std::ostream &out, tensor<T, n...> A)
+{
+   out << '{' << A[0];
+   for (int i = 1; i < tensor<T, n...>::shape[0]; i++)
+   {
+      out << ", " << A[i];
+   }
+   out << '}';
+   return out;
+}
+
+auto derivative_wrt(double a)
+{
+   return dual<double>{a, 1};
+}
+
+template<typename T, int m>
+auto derivative_wrt(tensor<T, m> A)
+{
+   tensor<dual<tensor<double, m>>, m> A_dual{};
+   for (int i = 0; i < m; i++)
+   {
+      A_dual[i].value = A[i];
+      A_dual[i].gradient[i] = 1.0;
+   }
+   return A_dual;
+}
+
+template<typename T, int m, int n>
+auto derivative_wrt(tensor<T, m, n> A)
+{
+   tensor<dual<tensor<double, m, n>>, m, n> A_dual{};
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         A_dual[i][j].value = A[i][j];
+         A_dual[i][j].gradient[i][j] = 1.0;
+      }
+   }
+   return A_dual;
+}
+
+template<typename grad_type, int nrows, int ncols>
+auto directional_derivative(tensor<dual<grad_type>, nrows, ncols> A, grad_type n)
+{
+   tensor<double, nrows, ncols> dA_dn{};
+   for (int i = 0; i < nrows; i++)
+   {
+      for (int j = 0; j < ncols; j++)
+      {
+         dA_dn[i][j] = inner(A[i][j].gradient, n);
+      }
+   }
+   return dA_dn;
+}

miniapps/variationalform/test_tensor_ad.cpp

@@ -0,0 +1,69 @@
+#include "mfem.hpp"
+#include "tensor.hpp"
+
+#define CATCH_CONFIG_MAIN
+#include "catch.hpp"
+
+using namespace mfem;
+
+TEST_CASE("Tensor tests", "[tensor]")
+{
+   SECTION("norm")
+   {
+      tensor<double, 5> a = {{1.0, 2.0, 3.0, 4.0, 5.0}};
+      REQUIRE(norm(a) - sqrt(55) == Approx(0.0));
+   }
+}
+
+TEST_CASE("Dual number tensor tests", "[DualNumber]")
+{
+   auto eps = std::numeric_limits<double>::epsilon();
+   double x = 0.5;
+
+   SECTION("cos")
+   {
+      auto xd = cos(derivative_wrt(x));
+      REQUIRE(abs(-sin(x) - xd.gradient) == Approx(0.0));
+   }
+
+   SECTION("exp")
+   {
+      auto xd = exp(derivative_wrt(x));
+      REQUIRE(abs(exp(x) - xd.gradient) == Approx(0.0));
+   }
+
+   SECTION("log")
+   {
+      auto xd = log(derivative_wrt(x));
+      REQUIRE(abs(1.0 / x - xd.gradient) == Approx(0.0));
+   }
+
+   SECTION("pow")
+   {
+      // f(x) = x^3/2
+      auto xd = pow(derivative_wrt(x), 1.5);
+      REQUIRE(abs(1.5 * pow(x, 0.5) - xd.gradient) == Approx(0.0));
+   }
+
+   SECTION("mixed operations")
+   {
+      auto xd = derivative_wrt(x);
+      auto r = cos(xd) * cos(xd);
+      REQUIRE(abs(-2.0 * sin(x) * cos(x) - r.gradient) == Approx(0.0));
+
+      r = exp(xd) * cos(xd);
+      REQUIRE(abs(exp(x) * (cos(x) - sin(x)) - r.gradient) < eps);
+
+      r = log(xd) * cos(xd);
+      REQUIRE(abs((cos(x) / x - log(x) * sin(x)) - r.gradient) < eps);
+
+      r = exp(xd) * pow(xd, 1.5);
+      REQUIRE(abs((exp(x) * (pow(x, 1.5) + 1.5 * pow(x, 0.5))) - r.gradient)
+              < eps);
+
+      tensor<double, 2> vx = {{0.5, 0.25}};
+      tensor<double, 2> vre = {{0.894427190999916, 0.4472135954999579}};
+      auto vr = norm(derivative_wrt(vx));
+      REQUIRE(norm(vr.gradient - vre) < eps);
+   }
+}