Adjust the intructions just a bit more

Better for testing vjp, since strain field is not uniform.

CMakeLists.txt

@@ -522,7 +522,10 @@ endif()
 
 # Enzyme
 if (MFEM_USE_ENZYME)
-   find_package(ENZYME REQUIRED)
+    find_package(Enzyme REQUIRED HINTS ${ENZYME_DIR})
+    message(STATUS "Enzyme found in ${ENZYME_DIR}.")
+    set(ENZYME_INCLUDE_DIRS ${ENZYME_DIR}/include)
+    set(ENZYME_FOUND 1)
 endif()
 
 # MFEM_TIMER_TYPE
@@ -629,6 +632,11 @@ set(MFEM_INSTALL_DIR ${CMAKE_INSTALL_PREFIX} CACHE PATH
 mfem_add_library(mfem ${SOURCES} ${HEADERS} ${MASTER_HEADERS})
 # message(STATUS "TPL_LIBRARIES = ${TPL_LIBRARIES}")
 target_link_libraries(mfem PUBLIC ${TPL_LIBRARIES})
+
+if (MFEM_USE_ENZYME)
+  target_link_libraries(mfem PUBLIC ClangEnzymeFlags)
+endif()
+
 if (MINGW)
   target_link_libraries(mfem PRIVATE ws2_32)
 endif()

config/cmake/modules/FindENZYME.cmake

@@ -1,27 +0,0 @@
-# Copyright (c) 2010-2024, Lawrence Livermore National Security, LLC. Produced
-# at the Lawrence Livermore National Laboratory. All Rights reserved. See files
-# LICENSE and NOTICE for details. LLNL-CODE-806117.
-#
-# This file is part of the MFEM library. For more information and source code
-# availability visit https://mfem.org.
-#
-# MFEM is free software; you can redistribute it and/or modify it under the
-# terms of the BSD-3 license. We welcome feedback and contributions, see file
-# CONTRIBUTING.md for details.
-
-message(STATUS "Looking for ENZYME ...")
-message(STATUS "   in ENZYME_DIR = ${ENZYME_DIR}")
-
-# Make sure the directory and version combination works. Do nothing otherwise.
-if(EXISTS "${ENZYME_DIR}/ClangEnzyme-${ENZYME_VERSION}.so")
-  message(STATUS "Found ENZYME: ${ENZYME_DIR}/ClangEnzyme-${ENZYME_VERSION}.so")
-
-  # Set ENZYME_FOUND
-  set(ENZYME_FOUND TRUE CACHE BOOL "ENZYME was found." FORCE)
-
-  # Set CXX flags to accommodate the Enzyme Clang plugin
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Xclang -load -Xclang ${ENZYME_DIR}/ClangEnzyme-${ENZYME_VERSION}.so -mllvm -enzyme-loose-types=1")
-  set(MFEM_USE_ENZYME YES)
-else()
-
-endif()

data/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 0 1
+2 1 1 2
+3 1 2 3
+4 1 3 0
+
+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

data/skewed-square.mesh

@@ -0,0 +1,35 @@
+MFEM mesh v1.0
+
+#
+# MFEM Geometry Types (see mesh/geom.hpp):
+#
+# POINT       = 0
+# SEGMENT     = 1
+# TRIANGLE    = 2
+# SQUARE      = 3
+# TETRAHEDRON = 4
+# CUBE        = 5
+# PRISM       = 6
+#
+
+dimension
+2
+
+elements
+1
+1 3 0 1 2 3
+
+boundary
+4
+1 1 0 1
+2 1 1 2
+3 1 2 3
+4 1 3 0
+
+vertices
+4
+2
+0 0
+1 0.3
+1.4 1.2
+0.25 1.34

examples/CMakeLists.txt

@@ -50,6 +50,41 @@ list(APPEND ALL_EXE_SRCS
 
 if (MFEM_USE_MPI)
   list(APPEND ALL_EXE_SRCS
+    dfem_poisson.cpp
+    enzyme_interface_smoketest.cpp
+    test_dfem_dual.cpp
+    test_dfem.cpp
+    dfem_laghos.cpp
+    dfem_minimal_example.cpp
+    dfem_test_diffusion_2d.cpp
+    dfem_test_diffusion.cpp
+    dfem_test_ordering.cpp
+    dfem_test_vector_diffusion.cpp
+    dfem_test_elasticity.cpp
+    dfem_test_nonlinear_elasticity_3d.cpp
+    dfem_test_nonlinear_diffusion_3d.cpp
+    dfem_test_interpolate_linear_scalar.cpp
+    dfem_test_interpolate_linear_scalar_3d.cpp
+    dfem_test_interpolate_gradient_linear_scalar_3d.cpp
+    dfem_test_interpolate_gradient_linear_scalar.cpp
+    dfem_test_mass_scalar_3d.cpp
+    dfem_test_mass_scalar_2d.cpp
+    dfem_test_interpolate_linear_vector.cpp
+    dfem_test_interpolate_linear_vector_3d.cpp
+    dfem_test_objective_vjp.cpp
+    dfem_test_objective.cpp
+    dfem_nonlinear_advdiff.cpp
+    dfem_advection_supg.cpp
+    dfem_navier_stokes.cpp
+    dfem_stokes.cpp
+    dfem_nonlinear_elasticity.cpp
+    dfem_fsi.cpp
+    dfem_cfd.cpp
+    dfem_heat.cpp
+    dfem_navier_stokes_edac.cpp
+    dfem_patch_test
+    dfem_plasticity.cpp
+    dfem_elasticity_vjp.cpp
     ex0p.cpp
     ex1p.cpp
     ex2p.cpp
@@ -110,6 +145,29 @@ include_directories(BEFORE ${PROJECT_BINARY_DIR})
 # Add one executable per cpp file
 add_mfem_examples(ALL_EXE_SRCS)
 
+target_link_libraries(dfem_poisson ClangEnzymeFlags)
+target_link_libraries(dfem_stokes ClangEnzymeFlags)
+target_link_libraries(enzyme_interface_smoketest ClangEnzymeFlags)
+target_link_libraries(test_dfem ClangEnzymeFlags)
+target_link_libraries(dfem_laghos ClangEnzymeFlags)
+target_link_libraries(dfem_minimal_example ClangEnzymeFlags)
+target_link_libraries(dfem_test_diffusion ClangEnzymeFlags)
+target_link_libraries(dfem_test_nonlinear_diffusion_3d ClangEnzymeFlags)
+target_link_libraries(dfem_test_nonlinear_elasticity_3d ClangEnzymeFlags)
+target_link_libraries(dfem_test_objective_vjp ClangEnzymeFlags)
+target_link_libraries(dfem_test_objective ClangEnzymeFlags)
+target_link_libraries(dfem_nonlinear_advdiff ClangEnzymeFlags)
+target_link_libraries(dfem_navier_stokes ClangEnzymeFlags)
+target_link_libraries(dfem_navier_stokes_edac ClangEnzymeFlags)
+target_link_libraries(dfem_advection_supg ClangEnzymeFlags)
+target_link_libraries(dfem_nonlinear_elasticity ClangEnzymeFlags)
+target_link_libraries(dfem_fsi ClangEnzymeFlags)
+target_link_libraries(dfem_heat ClangEnzymeFlags)
+target_link_libraries(dfem_cfd ClangEnzymeFlags)
+target_link_libraries(dfem_patch_test ClangEnzymeFlags)
+target_link_libraries(dfem_plasticity ClangEnzymeFlags)
+target_link_libraries(dfem_elasticity_vjp ClangEnzymeFlags)
+
 # Add a test for each example
 if (MFEM_ENABLE_TESTING)
   foreach(SRC_FILE ${ALL_EXE_SRCS})

examples/dfem/cxxtimer.hpp

@@ -0,0 +1,184 @@
+/*
+
+MIT License
+
+Copyright (c) 2017 André L. Maravilha
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+*/
+
+#ifndef CXX_TIMER_HPP
+#define CXX_TIMER_HPP
+
+#include <chrono>
+
+
+namespace cxxtimer {
+
+/**
+ * This class works as a stopwatch.
+ */
+class Timer {
+
+public:
+
+    /**
+     * Constructor.
+     *
+     * @param   start
+     *          If true, the timer is started just after construction.
+     *          Otherwise, it will not be automatically started.
+     */
+    Timer(bool start = false);
+
+    /**
+     * Copy constructor.
+     *
+     * @param   other
+     *          The object to be copied.
+     */
+    Timer(const Timer& other) = default;
+
+    /**
+     * Transfer constructor.
+     *
+     * @param   other
+     *          The object to be transferred.
+     */
+    Timer(Timer&& other) = default;
+
+    /**
+     * Destructor.
+     */
+    virtual ~Timer() = default;
+
+    /**
+     * Assignment operator by copy.
+     *
+     * @param   other
+     *          The object to be copied.
+     *
+     * @return  A reference to this object.
+     */
+    Timer& operator=(const Timer& other) = default;
+
+    /**
+     * Assignment operator by transfer.
+     *
+     * @param   other
+     *          The object to be transferred.
+     *
+     * @return  A reference to this object.
+     */
+    Timer& operator=(Timer&& other) = default;
+
+    /**
+     * Start/resume the timer.
+     */
+    void start();
+
+    /**
+     * Stop/pause the timer.
+     */
+    void stop();
+
+    /**
+     * Reset the timer.
+     */
+    void reset();
+
+    /**
+     * Return the elapsed time.
+     *
+     * @param   duration_t
+     *          The duration type used to return the time elapsed. If not
+     *          specified, it returns the time as represented by
+     *          std::chrono::milliseconds.
+     *
+     * @return  The elapsed time.
+     */
+    template <class duration_t = std::chrono::milliseconds>
+    typename duration_t::rep count() const;
+
+private:
+
+    bool started_;
+    bool paused_;
+    std::chrono::steady_clock::time_point reference_;
+    std::chrono::duration<long double> accumulated_;
+};
+
+}
+
+
+inline cxxtimer::Timer::Timer(bool start) :
+        started_(false), paused_(false),
+        reference_(std::chrono::steady_clock::now()),
+        accumulated_(std::chrono::duration<long double>(0)) {
+    if (start) {
+        this->start();
+    }
+}
+
+inline void cxxtimer::Timer::start() {
+    if (!started_) {
+        started_ = true;
+        paused_ = false;
+        accumulated_ = std::chrono::duration<long double>(0);
+        reference_ = std::chrono::steady_clock::now();
+    } else if (paused_) {
+        reference_ = std::chrono::steady_clock::now();
+        paused_ = false;
+    }
+}
+
+inline void cxxtimer::Timer::stop() {
+    if (started_ && !paused_) {
+        std::chrono::steady_clock::time_point now = std::chrono::steady_clock::now();
+        accumulated_ = accumulated_ + std::chrono::duration_cast< std::chrono::duration<long double> >(now - reference_);
+        paused_ = true;
+    }
+}
+
+inline void cxxtimer::Timer::reset() {
+    if (started_) {
+        started_ = false;
+        paused_ = false;
+        reference_ = std::chrono::steady_clock::now();
+        accumulated_ = std::chrono::duration<long double>(0);
+    }
+}
+
+template <class duration_t>
+typename duration_t::rep cxxtimer::Timer::count() const {
+    if (started_) {
+        if (paused_) {
+            return std::chrono::duration_cast<duration_t>(accumulated_).count();
+        } else {
+            return std::chrono::duration_cast<duration_t>(
+                    accumulated_ + (std::chrono::steady_clock::now() - reference_)).count();
+        }
+    } else {
+        return duration_t(0).count();
+    }
+}
+
+
+#endif

examples/dfem/dfem.hpp

@@ -0,0 +1,4 @@
+#pragma once
+
+#include "dfem_differentiable_operator.hpp"
+#include "dfem_element_operator.hpp"

examples/dfem/dfem_action_callback.icc

@@ -0,0 +1,232 @@
+template <
+   typename kernels_tuple,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t num_fields,
+   size_t num_kernels
+   >
+template <
+   typename kernel_t
+   >
+void DifferentiableOperator<kernels_tuple,
+     num_solutions,
+     num_parameters,
+     num_fields,
+     num_kernels>::Action::create_action_callback(
+        kernel_t kernel,
+        mult_func_t &func)
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                     kernel.inputs, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                      kernel.outputs, std::make_index_sequence<kernel.num_koutputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(op.fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(kernel.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
+   const int num_entities = GetNumEntities<entity_t>(op.mesh);
+   const int num_qp = op.integration_rule.GetNPoints();
+
+   // All solutions T-vector sizes make up the width of the operator, since
+   // they are explicitly provided in Mult() for example.
+
+   op.width = GetTrueVSize(op.fields[test_space_field_idx]);
+   op.residual_lsize = GetVSize(op.fields[test_space_field_idx]);
+
+   if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      op.height = 1;
+   }
+   else
+   {
+      op.height = op.residual_lsize;
+   }
+
+   residual_l.SetSize(op.residual_lsize);
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : op.fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+                                              doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/op.mesh.Dimension()) + 0.5);
+
+   residual_e.SetSize(R->Height());
+
+   const int residual_size_on_qp = GetSizeOnQP<entity_t>(
+                                      mfem::get<hardcoded_output_idx>(kernel.outputs),
+                                      op.fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(kernel.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(kernel.outputs, dtq,
+                                                    koutput_to_field);
+
+   auto input_fops = create_bare_fops(kernel.inputs);
+   auto output_fops = create_bare_fops(kernel.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(output_fops).size_on_qp /
+      mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(output_fops).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(this->op.integration_rule.GetWeights().Read(),
+                             num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
+                                                std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto shmem_info = get_shmem_info<entity_t>(input_dtq_maps,
+                                              output_dtq_maps,
+                                              op.fields,
+                                              num_entities,
+                                              kernel.inputs,
+                                              num_qp,
+                                              input_size_on_qp,
+                                              residual_size_on_qp);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   print_shared_memory_info(shmem_info);
+
+   func = [=](Vector &ye_mem) mutable
+   {
+      restriction<entity_t>(op.solutions, solutions_l, this->fields_e,
+                            op.element_dof_ordering);
+      restriction<entity_t>(op.parameters, parameters_l, this->fields_e,
+                            op.element_dof_ordering,
+                            op.solutions.size());
+
+      auto ye = Reshape(ye_mem.ReadWrite(), test_vdim, num_test_dof, num_entities);
+      auto wrapped_fields_e = wrap_fields(this->fields_e, shmem_info.field_sizes, num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, void *shmem)
+      {
+         // printf("\ne: %d\n", e);
+         // tic();
+         auto input_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT_DTQ],
+            shmem_info.input_dtq_sizes,
+            input_dtq_maps);
+
+         auto output_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT_DTQ],
+            shmem_info.output_dtq_sizes,
+            output_dtq_maps);
+
+         auto fields_shmem = load_field_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::FIELD],
+            shmem_info.field_sizes,
+            kinput_to_field,
+            wrapped_fields_e,
+            e);
+
+         // These methods don't copy, they simply create a `DeviceTensor` object
+         // that points to correct chunks of the shared memory pool.
+         auto input_shmem = load_input_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT],
+            shmem_info.input_sizes,
+            num_qp);
+
+         auto residual_shmem = load_residual_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT],
+            shmem_info.residual_size,
+            num_qp);
+
+         auto scratch_mem = load_scratch_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::TEMP],
+            shmem_info.temp_sizes);
+
+         MFEM_SYNC_THREAD;
+         // printf("shmem load elapsed: %.1fus\n", toc() * 1e6);
+
+         // tic();
+         map_fields_to_quadrature_data<TensorProduct>(
+            input_shmem, fields_shmem, input_dtq_shmem, input_fops, ir_weights, scratch_mem,
+         std::make_index_sequence<kernel.num_kinputs> {});
+         // printf("interpolate elapsed: %.1fus\n", toc() * 1e6);
+
+         // tic();
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  const int q = qx + q1d * (qy + q1d * qz);
+                  auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+                  auto r = Reshape(&residual_shmem(0, q), residual_size_on_qp);
+                  apply_kernel(r, kernel.func, kernel_args, input_shmem, q);
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+         // printf("qf elapsed: %.1fus\n", toc() * 1e6);
+
+         // tic();
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                      mfem::get<0>(output_fops),
+                                                      output_dtq_shmem[hardcoded_output_idx],
+                                                      scratch_mem);
+         // printf("integrate elapsed: %.1fus\n", toc() * 1e6);
+
+      }, num_entities, q1d, q1d, q1d, shmem_info.total_size, shmem_cache.ReadWrite());
+
+      if constexpr (std::is_same_v<decltype(output_fop), None>)
+      {
+         residual_l = ye_mem;
+      }
+      else
+      {
+         R->MultTranspose(ye_mem, residual_l);
+      }
+   };
+
+   if constexpr (std::is_same_v<decltype(output_fop), None>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         y = r_local;
+      };
+   }
+   else if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         double local_sum = r_local.Sum();
+         MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM,
+                       op.mesh.GetComm());
+         MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      };
+   }
+   else
+   {
+      auto P = get_prolongation(op.fields[test_space_field_idx]);
+      prolongation_transpose = [P](const Vector &r_local, Vector &y)
+      {
+         P->MultTranspose(r_local, y);
+      };
+   }
+}

examples/dfem/dfem_assemble_hypreparmatrix.icc

@@ -0,0 +1,308 @@
+template <
+   typename kernels_tuple,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t num_fields,
+   size_t num_kernels
+   >
+template <
+   size_t derivative_idx
+   >
+template <
+   typename kernel_t
+   >
+void DifferentiableOperator<kernels_tuple,
+     num_solutions,
+     num_parameters,
+     num_fields,
+     num_kernels>::Derivative<derivative_idx>::assemble_hypreparmatrix_impl(
+        kernel_t kernel, HypreParMatrix &A)
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                     kernel.inputs,
+                                                                     std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                      kernel.outputs,
+                                                                      std::make_index_sequence<kernel.num_koutputs> {});
+
+   auto output_fop = std::get<0>(kernel.outputs);
+
+   constexpr int hardcoded_output_idx = 0;
+
+   int num_qp = op.integration_rule.GetNPoints();;
+   int num_el = 0;
+   int dimension = 0;
+   if constexpr (std::is_same_v<entity_t, Entity::Element>)
+   {
+      num_el = op.mesh.GetNE();
+      dimension = op.dim;
+   }
+   else if (std::is_same_v<entity_t, Entity::Face>)
+   {
+      num_el = op.mesh.GetNumFacesWithGhost();
+      dimension = op.dim - 1;
+   }
+   else
+   {
+      static_assert(always_false<entity_t>, "not implemented");
+   }
+
+   std::vector<const DofToQuad*> dtqmaps;
+   for (const auto &field : op.fields)
+   {
+      dtqmaps.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+                                                  doftoquad_mode));
+   }
+
+   // Allocate memory for fields on quadrature points
+   auto input_qp_mem = create_input_qp_memory(num_qp, kernel.inputs,
+                                              std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto directions_qp_mem = create_input_qp_memory(num_qp, kernel.inputs,
+                                                   std::make_index_sequence<kernel.num_kinputs> {});
+
+   for (auto &d_qp_mem : directions_qp_mem)
+   {
+      d_qp_mem = 0.0;
+   }
+
+   std::array<bool, kernel.num_kinputs> kinput_is_dependent;
+   bool no_kinput_is_dependent = true;
+   for (int i = 0; i < kinput_is_dependent.size(); i++)
+   {
+      if (kinput_to_field[i] == derivative_idx)
+      {
+         no_kinput_is_dependent = false;
+         kinput_is_dependent[i] = true;
+         // out << "function input " << i << " is dependent on "
+         //     << op.fields[kinput_to_field[i]].field_label << "\n";
+      }
+      else
+      {
+         kinput_is_dependent[i] = false;
+      }
+   }
+
+   if (no_kinput_is_dependent)
+   {
+      return;
+   }
+
+   auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+   auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+
+   DeviceTensor<1, const double> integration_weights(
+      this->op.integration_rule.GetWeights().Read(), num_qp);
+
+   Vector zero;
+   GeometricFactorMaps geometric_factors
+   {
+      DeviceTensor<3, const double>(zero.Read(), 0, 0, 0)
+   };
+
+   // fields interpolated to the quadrature points in the order of
+   // kernel function arguments
+   auto input_qp = map_inputs_to_memory(input_qp_mem, num_qp,
+                                        kernel.inputs,
+                                        std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto directions_qp = map_inputs_to_memory(directions_qp_mem, num_qp,
+                                             kernel.inputs,
+                                             std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto input_dtq_ops = create_dtq_operators<entity_t>(kernel.inputs, dtqmaps,
+                                                       kinput_to_field);
+   auto dependent_input_dtq_ops = create_dtq_operators_conditional<entity_t>(
+                                     kernel.inputs,
+                                     dtqmaps,
+                                     kinput_to_field,
+                                     kinput_is_dependent, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto output_dtq_ops = create_dtq_operators<entity_t>(kernel.outputs, dtqmaps,
+                                                        koutput_to_field);
+
+   constexpr int fixed_output_idx = 0;
+   auto Bv = output_dtq_ops[fixed_output_idx];
+   auto [num_test_qp, test_op_dim, num_test_dof] = Bv.GetShape();
+   const int test_vdim = std::get<0>(kernel.outputs).vdim;
+
+   const int num_trial_dof = dependent_input_dtq_ops[0].GetShape()[2];
+   int trial_vdim = 0;
+   for (int i = 0; i < kinput_is_dependent.size(); i++)
+   {
+      if (kinput_is_dependent[i])
+      {
+         trial_vdim = GetVDim(op.fields[kinput_to_field[i]]);
+         break;
+      }
+   }
+
+   // All trial operators dimensions accumulated
+   int total_trial_op_dim = 0;
+   for (int s = 0; s < dependent_input_dtq_ops.size(); s++)
+   {
+      total_trial_op_dim += dependent_input_dtq_ops[s].GetShape()[1];
+   }
+
+   Vector a_qp_mem(test_vdim * test_op_dim * trial_vdim * total_trial_op_dim *
+                   num_qp *
+                   num_el);
+   const auto a_qp = Reshape(a_qp_mem.ReadWrite(), test_vdim, test_op_dim,
+                             trial_vdim, total_trial_op_dim, num_qp,
+                             num_el);
+
+   Vector Ae_mem(num_test_dof * test_vdim * num_trial_dof * trial_vdim * num_el);
+   Ae_mem = 0.0;
+
+   auto A_e = Reshape(Ae_mem.ReadWrite(), num_test_dof, test_vdim, num_trial_dof,
+                      trial_vdim, num_el);
+
+   for (int e = 0; e < num_el; e++)
+   {
+      map_fields_to_quadrature_data(
+         input_qp, e, this->fields_e,
+         kinput_to_field, input_dtq_ops,
+         integration_weights, geometric_factors, kernel.inputs,
+         std::make_index_sequence<kernel.num_kinputs> {});
+
+      for (int q = 0; q < num_qp; q++)
+      {
+         for (int j = 0; j < trial_vdim; j++)
+         {
+            size_t m_offset = 0;
+            for (int s = 0; s < dependent_input_dtq_ops.size(); s++)
+            {
+               auto Bu = dependent_input_dtq_ops[s];
+               auto [unused1, trial_op_dim, unused2] = Bu.GetShape();
+               auto d_qp = Reshape(&(directions_qp[Bu.which_input])[0], trial_vdim,
+                                   trial_op_dim, num_qp);
+               for (int m = 0; m < trial_op_dim; m++)
+               {
+                  d_qp(j, m, q) = 1.0;
+                  Vector f_qp = apply_kernel_fwddiff_enzyme(
+                                   kernel.func,
+                                   kernel_args,
+                                   input_qp,
+                                   kernel_shadow_args,
+                                   directions_qp,
+                                   q);
+                  // Vector f_qp = apply_kernel_fwddiff_dual(
+                  //                  kernel.func,
+                  //                  kernel_args,
+                  //                  input_qp,
+                  //                  directions_qp,
+                  //                  q);
+                  d_qp(j, m, q) = 0.0;
+
+                  auto f = Reshape(f_qp.Read(), test_vdim, test_op_dim);
+
+                  for (int i = 0; i < test_vdim; i++)
+                  {
+                     for (int k = 0; k < test_op_dim; k++)
+                     {
+                        a_qp(i, k, j, m + m_offset, q, e) = f(i, k);
+                     }
+                  }
+               }
+               m_offset += trial_op_dim;
+            }
+         }
+      }
+
+      Vector fhat_mem(test_op_dim * num_qp * dimension);
+      auto fhat = Reshape(fhat_mem.ReadWrite(), test_vdim, test_op_dim, num_qp);
+      for (int J = 0; J < num_trial_dof; J++)
+      {
+         for (int j = 0; j < trial_vdim; j++)
+         {
+            fhat_mem = 0.0;
+            size_t m_offset = 0;
+            for (int s = 0; s < dependent_input_dtq_ops.size(); s++)
+            {
+               auto Bu = dependent_input_dtq_ops[s];
+               int trial_op_dim = dependent_input_dtq_ops[s].GetShape()[1];
+               for (int q = 0; q < num_qp; q++)
+               {
+                  for (int i = 0; i < test_vdim; i++)
+                  {
+                     for (int k = 0; k < test_op_dim; k++)
+                     {
+                        for (int m = 0; m < trial_op_dim; m++)
+                        {
+                           fhat(i, k, q) += a_qp(i, k, j, m + m_offset, q, e) * Bu(q, m, J);
+                        }
+                     }
+                  }
+               }
+               m_offset += trial_op_dim;
+            }
+
+            auto bvtfhat = Reshape(&A_e(0, 0, J, j, e), num_test_dof, test_vdim);
+            map_quadrature_data_to_fields(bvtfhat, fhat, output_fop,
+                                          output_dtq_ops[hardcoded_output_idx]);
+         }
+      }
+   }
+
+   bool same_test_and_trial = false;
+   if (koutput_to_field[0] ==
+       kinput_to_field[dependent_input_dtq_ops[0].which_input])
+   {
+      same_test_and_trial = true;
+   }
+
+   auto trial_fes = *std::get_if<const ParFiniteElementSpace *>
+                    (&op.fields[kinput_to_field[dependent_input_dtq_ops[0].which_input]].data);
+
+   auto test_fes = *std::get_if<const ParFiniteElementSpace *>
+                   (&op.fields[koutput_to_field[0]].data);
+
+   SparseMatrix mat(test_fes->GlobalVSize(), trial_fes->GlobalVSize());
+
+   if (test_fes == nullptr)
+   {
+      MFEM_ABORT("error");
+   }
+
+   for (int e = 0; e < num_el; e++)
+   {
+      auto tmp = Reshape(Ae_mem.ReadWrite(), num_test_dof * test_vdim,
+                         num_trial_dof * trial_vdim,
+                         num_el);
+      DenseMatrix A_e(&tmp(0, 0, e), num_test_dof * test_vdim,
+                      num_trial_dof * trial_vdim);
+      Array<int> test_vdofs, trial_vdofs;
+      test_fes->GetElementVDofs(e, test_vdofs);
+      GetElementVDofs(
+         op.fields[kinput_to_field[dependent_input_dtq_ops[0].which_input]], e,
+         trial_vdofs);
+      mat.AddSubMatrix(test_vdofs, trial_vdofs, A_e, 1);
+   }
+   mat.Finalize();
+
+   if (same_test_and_trial)
+   {
+      HypreParMatrix tmp(test_fes->GetComm(),
+                         test_fes->GlobalVSize(),
+                         test_fes->GetDofOffsets(),
+                         &mat);
+
+      A = *RAP(&tmp, test_fes->Dof_TrueDof_Matrix());
+      A.EliminateBC(op.ess_tdof_list, DiagonalPolicy::DIAG_ONE);
+   }
+   else
+   {
+      HypreParMatrix tmp(test_fes->GetComm(),
+                         test_fes->GlobalVSize(),
+                         trial_fes->GlobalVSize(),
+                         test_fes->GetDofOffsets(),
+                         trial_fes->GetDofOffsets(),
+                         &mat);
+
+      A = *RAP(test_fes->Dof_TrueDof_Matrix(), &tmp, trial_fes->Dof_TrueDof_Matrix());
+      // A.EliminateBC(op.ess_tdof_list, DiagonalPolicy::DIAG_ONE);
+   }
+}

examples/dfem/dfem_assemble_vector.icc

@@ -0,0 +1,233 @@
+template <
+   typename kernels_tuple,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t num_fields,
+   size_t num_kernels
+   >
+template <
+   size_t derivative_idx
+   >
+template <
+   typename kernel_t
+   >
+void DifferentiableOperator<kernels_tuple,
+     num_solutions,
+     num_parameters,
+     num_fields,
+     num_kernels>::Derivative<derivative_idx>::assemble_vector_impl(
+        kernel_t kernel, Vector &v)
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                     kernel.inputs,
+                                                                     std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                      kernel.outputs,
+                                                                      std::make_index_sequence<kernel.num_koutputs> {});
+
+   auto output_fop = std::get<0>(kernel.outputs);
+
+   constexpr int hardcoded_output_idx = 0;
+
+   int num_qp = op.integration_rule.GetNPoints();;
+   int num_el = 0;
+   int dimension = 0;
+   if constexpr (std::is_same_v<entity_t, Entity::Element>)
+   {
+      num_el = op.mesh.GetNE();
+      dimension = op.dim;
+   }
+   else if (std::is_same_v<entity_t, Entity::Face>)
+   {
+      num_el = op.mesh.GetNumFacesWithGhost();
+      dimension = op.dim - 1;
+   }
+   else
+   {
+      static_assert(always_false<entity_t>, "not implemented");
+   }
+
+   std::vector<const DofToQuad*> dtqmaps;
+   for (const auto &field : op.fields)
+   {
+      dtqmaps.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+                                                  doftoquad_mode));
+   }
+
+   // Allocate memory for fields on quadrature points
+   auto input_qp_mem = create_input_qp_memory(num_qp, kernel.inputs,
+                                              std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto directions_qp_mem = create_input_qp_memory(num_qp, kernel.inputs,
+                                                   std::make_index_sequence<kernel.num_kinputs> {});
+
+   for (auto &d_qp_mem : directions_qp_mem)
+   {
+      d_qp_mem = 0.0;
+   }
+
+   std::array<bool, kernel.num_kinputs> kinput_is_dependent;
+   bool no_kinput_is_dependent = true;
+   for (int i = 0; i < kinput_is_dependent.size(); i++)
+   {
+      if (kinput_to_field[i] == derivative_idx)
+      {
+         no_kinput_is_dependent = false;
+         kinput_is_dependent[i] = true;
+         // out << "function input " << i << " is dependent on "
+         //     << op.fields[kinput_to_field[i]].field_label << "\n";
+      }
+      else
+      {
+         kinput_is_dependent[i] = false;
+      }
+   }
+
+   if (no_kinput_is_dependent)
+   {
+      return;
+   }
+
+   auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+   auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+
+   DeviceTensor<1, const double> integration_weights(
+      this->op.integration_rule.GetWeights().Read(), num_qp);
+
+   Vector zero;
+   GeometricFactorMaps geometric_factors
+   {
+      DeviceTensor<3, const double>(zero.Read(), 0, 0, 0)
+   };
+
+   // fields interpolated to the quadrature points in the order of
+   // kernel function arguments
+   auto input_qp = map_inputs_to_memory(input_qp_mem, num_qp,
+                                        kernel.inputs,
+                                        std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto directions_qp = map_inputs_to_memory(directions_qp_mem, num_qp,
+                                             kernel.inputs,
+                                             std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto input_dtq_ops = create_dtq_operators<entity_t>(kernel.inputs, dtqmaps,
+                                                       kinput_to_field);
+   auto dependent_input_dtq_ops = create_dtq_operators_conditional<entity_t>(
+                                     kernel.inputs,
+                                     dtqmaps,
+                                     kinput_to_field,
+                                     kinput_is_dependent, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto output_dtq_ops = create_dtq_operators<entity_t>(kernel.outputs, dtqmaps,
+                                                        koutput_to_field);
+
+   constexpr int fixed_output_idx = 0;
+   auto Bv = output_dtq_ops[fixed_output_idx];
+   auto [num_test_qp, test_op_dim, num_test_dof] = Bv.GetShape();
+   const int test_vdim = std::get<0>(kernel.outputs).vdim;
+
+   const int num_trial_dof = dependent_input_dtq_ops[0].GetShape()[2];
+   int trial_vdim = 0;
+   int dependent_field_idx = -1;
+   for (int i = 0; i < kinput_is_dependent.size(); i++)
+   {
+      if (kinput_is_dependent[i])
+      {
+         dependent_field_idx = kinput_to_field[i];
+         break;
+      }
+   }
+
+   trial_vdim = GetVDim(op.fields[dependent_field_idx]);
+
+   // All trial operators dimensions accumulated
+   int total_trial_op_dim = 0;
+   for (int s = 0; s < dependent_input_dtq_ops.size(); s++)
+   {
+      total_trial_op_dim += dependent_input_dtq_ops[s].GetShape()[1];
+   }
+
+   Vector a_qp_mem(trial_vdim * total_trial_op_dim * num_qp * num_el);
+   const auto a_qp = Reshape(a_qp_mem.ReadWrite(), trial_vdim,
+                             total_trial_op_dim, num_qp, num_el);
+   Vector ve_mem(num_trial_dof * trial_vdim * num_el);
+   ve_mem = 0.0;
+
+   for (int e = 0; e < num_el; e++)
+   {
+      map_fields_to_quadrature_data(
+         input_qp, e, this->fields_e,
+         kinput_to_field, input_dtq_ops,
+         integration_weights, geometric_factors, kernel.inputs,
+         std::make_index_sequence<kernel.num_kinputs> {});
+
+      for (int q = 0; q < num_qp; q++)
+      {
+         for (int j = 0; j < trial_vdim; j++)
+         {
+            size_t m_offset = 0;
+            for (int s = 0; s < dependent_input_dtq_ops.size(); s++)
+            {
+               auto Bu = dependent_input_dtq_ops[s];
+               auto [unused1, trial_op_dim, unused2] = Bu.GetShape();
+               auto d_qp = Reshape(&(directions_qp[Bu.which_input])[0], trial_vdim,
+                                   trial_op_dim, num_qp);
+               for (int m = 0; m < trial_op_dim; m++)
+               {
+                  d_qp(j, m, q) = 1.0;
+                  // Vector f_qp = apply_kernel_fwddiff_dual(
+                  //                  kernel.func,
+                  //                  kernel_args,
+                  //                  input_qp,
+                  //                  directions_qp,
+                  //                  q);
+                  Vector f_qp = apply_kernel_fwddiff_enzyme(
+                                   kernel.func,
+                                   kernel_args,
+                                   input_qp,
+                                   kernel_shadow_args,
+                                   directions_qp,
+                                   q);
+                  d_qp(j, m, q) = 0.0;
+
+                  auto f = Reshape(f_qp.Read(), test_vdim);
+                  a_qp(j, m + m_offset, q, e) = f(0);
+               }
+               m_offset += trial_op_dim;
+            }
+         }
+      }
+
+      auto shat = Reshape(ve_mem.ReadWrite(), num_trial_dof, trial_vdim, num_el);
+      for (int J = 0; J < num_trial_dof; J++)
+      {
+         for (int j = 0; j < trial_vdim; j++)
+         {
+            size_t m_offset = 0;
+            for (int s = 0; s < dependent_input_dtq_ops.size(); s++)
+            {
+               auto Bu = dependent_input_dtq_ops[s];
+               int trial_op_dim = dependent_input_dtq_ops[s].GetShape()[1];
+               for (int q = 0; q < num_qp; q++)
+               {
+                  for (int m = 0; m < trial_op_dim; m++)
+                  {
+                     shat(J, j, e) += a_qp(j, m + m_offset, q, e) * Bu(q, m, J);
+                  }
+               }
+               m_offset += trial_op_dim;
+            }
+         }
+      }
+   }
+
+   auto R = get_element_restriction(op.fields[dependent_field_idx],
+                                    element_dof_ordering);
+   Vector ve(R->Width());
+   R->MultTranspose(ve_mem, ve);
+
+   get_prolongation(op.fields[dependent_field_idx])->MultTranspose(ve, v);
+}

examples/dfem/dfem_derivative_callback.icc

@@ -0,0 +1,244 @@
+template <
+   typename kernels_tuple,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t num_fields,
+   size_t num_kernels
+   >
+template <
+   size_t derivative_idx
+   >
+template <
+   typename kernel_t
+   >
+void DifferentiableOperator<kernels_tuple,
+     num_solutions,
+     num_parameters,
+     num_fields,
+     num_kernels>::Derivative<derivative_idx>::create_callback(kernel_t kernel,
+                                                               mult_func_t &func)
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                     kernel.inputs, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                      kernel.outputs, std::make_index_sequence<kernel.num_koutputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(op.fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(kernel.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
+   const int num_entities = GetNumEntities<entity_t>(op.mesh);
+   const int num_qp = op.integration_rule.GetNPoints();
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : op.fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+                                              doftoquad_mode));
+   }
+   const int q1d = dtq[0]->nqpt;
+
+   derivative_action_e.SetSize(R->Height());
+
+   const int da_size_on_qp = GetSizeOnQP<entity_t>(
+                                mfem::get<hardcoded_output_idx>(kernel.outputs),
+                                op.fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(kernel.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(kernel.outputs, dtq,
+                                                    koutput_to_field);
+
+   auto input_fops = create_bare_fops(kernel.inputs);
+   auto output_fops = create_bare_fops(kernel.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(output_fops).size_on_qp /
+      mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(output_fops).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(this->op.integration_rule.GetWeights().Read(),
+                             num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
+                                                std::make_index_sequence<kernel.num_kinputs> {});
+
+   // Check which qf inputs are dependent on the dependent variable
+   std::array<bool, kernel.num_kinputs> kinput_is_dependent;
+   bool no_kinput_is_dependent = true;
+   for (int i = 0; i < kinput_is_dependent.size(); i++)
+   {
+      if (kinput_to_field[i] == derivative_idx)
+      {
+         no_kinput_is_dependent = false;
+         kinput_is_dependent[i] = true;
+         // out << "function input " << i << " is dependent on "
+         //     << op.fields[kinput_to_field[i]].field_label << "\n";
+      }
+      else
+      {
+         kinput_is_dependent[i] = false;
+      }
+   }
+
+   bool with_derivatives = true;
+   auto shmem_info = get_shmem_info<entity_t>(input_dtq_maps,
+                                              output_dtq_maps,
+                                              op.fields,
+                                              num_entities,
+                                              kernel.inputs,
+                                              num_qp,
+                                              input_size_on_qp,
+                                              da_size_on_qp,
+                                              derivative_idx);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   print_shared_memory_info(shmem_info);
+
+   func = [=](Vector &ye_mem) mutable
+   {
+      if (no_kinput_is_dependent)
+      {
+         return;
+      }
+
+      restriction<entity_t>(direction, direction_l, direction_e,
+                            op.element_dof_ordering);
+
+      auto ye = Reshape(ye_mem.ReadWrite(), num_test_dof, test_vdim, num_entities);
+      auto wrapped_fields_e = wrap_fields(this->fields_e, shmem_info.field_sizes, num_entities);
+      auto wrapped_direction_e = Reshape(direction_e.Read(), shmem_info.direction_size, num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, double *shmem)
+      {
+         auto input_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT_DTQ],
+            shmem_info.input_dtq_sizes,
+            input_dtq_maps);
+
+         auto output_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT_DTQ],
+            shmem_info.output_dtq_sizes,
+            output_dtq_maps);
+
+         auto fields_shmem = load_field_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::FIELD],
+            shmem_info.field_sizes,
+            kinput_to_field,
+            wrapped_fields_e,
+            e);
+
+         auto direction_shmem = load_direction_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::DIRECTION],
+            shmem_info.direction_size,
+            wrapped_direction_e,
+            e);
+
+         // These methods don't copy, they simply create a `DeviceTensor` object
+         // that points to correct chunks of the shared memory pool.
+         auto input_shmem = load_input_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT],
+            shmem_info.input_sizes,
+            num_qp);
+
+         auto shadow_shmem = load_input_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::SHADOW],
+            shmem_info.input_sizes,
+            num_qp);
+
+         auto residual_shmem = load_residual_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT],
+            shmem_info.residual_size,
+            num_qp);
+
+         auto scratch_mem = load_scratch_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::TEMP],
+            shmem_info.temp_sizes);
+
+         map_fields_to_quadrature_data<TensorProduct>(
+            input_shmem, fields_shmem, input_dtq_shmem, input_fops, ir_weights, scratch_mem,
+         std::make_index_sequence<kernel.num_kinputs> {});
+
+         zero_all(shadow_shmem);
+         map_direction_to_quadrature_data_conditional<TensorProduct>(
+            shadow_shmem, direction_shmem, input_dtq_shmem, input_fops, ir_weights,
+            scratch_mem, kinput_is_dependent,
+            std::make_index_sequence<kernel.num_kinputs> {});
+
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  const int q = qx + q1d * (qy + q1d * qz);
+
+                  auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+                  auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+
+                  auto r = Reshape(&residual_shmem(0, q), da_size_on_qp);
+                  apply_kernel_fwddiff_enzyme(
+                     r,
+                     kernel.func,
+                     kernel_args,
+                     input_shmem,
+                     kernel_shadow_args,
+                     shadow_shmem,
+                     q);
+                  // printf(">>>>> WARNING: AD DISABLED\n");
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                      mfem::get<0>(output_fops),
+                                                      output_dtq_shmem[hardcoded_output_idx],
+                                                      scratch_mem);
+      }, num_entities, q1d, q1d, 1, shmem_info.total_size, shmem_cache.ReadWrite());
+
+      R->MultTranspose(ye_mem, derivative_action_l);
+   };
+
+   if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         double local_sum = r_local.Sum();
+         MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM,
+                       op.mesh.GetComm());
+         MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      };
+   }
+   else
+   {
+      auto P = get_prolongation(op.fields[test_space_field_idx]);
+      prolongation_transpose = [P](const Vector &r_local, Vector &y)
+      {
+         P->MultTranspose(r_local, y);
+      };
+   }
+}

examples/dfem/dfem_differentiable_operator.hpp

@@ -0,0 +1,820 @@
+#pragma once
+
+#include <algorithm>
+#include <cstdlib>
+#include <functional>
+#include <iostream>
+#include <utility>
+#include <variant>
+#include <vector>
+#include <type_traits>
+#include <mfem.hpp>
+#include <type_traits>
+#include "dfem_fieldoperator.hpp"
+#include "dfem_parametricspace.hpp"
+#include "general/tic_toc.hpp"
+#include "tuple.hpp"
+#include <linalg/tensor.hpp>
+#include <enzyme/utils>
+#include <enzyme/enzyme>
+#include "dfem_util.hpp"
+#include "dfem_interpolate.hpp"
+#include "dfem_qfunction.hpp"
+#include "dfem_qfunction_dual.hpp"
+#include "dfem_integrate.hpp"
+
+namespace mfem
+{
+
+using mult_func_t = std::function<void(Vector &)>;
+
+template <
+   typename kernels_tuple,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t num_fields = num_solutions + num_parameters,
+   size_t num_kernels = mfem::tuple_size<kernels_tuple>::value,
+   typename autodiff_t = AutoDiff::NativeDualNumber
+   >
+class DifferentiableOperator : public Operator
+{
+public:
+   DifferentiableOperator(DifferentiableOperator&) = delete;
+   DifferentiableOperator(DifferentiableOperator&&) = delete;
+
+   class Action : public Operator
+   {
+   public:
+      template <typename kernel_t>
+      void create_action_callback(kernel_t kernel, mult_func_t &func);
+
+      template<std::size_t... idx>
+      void materialize_callbacks(kernels_tuple &ks,
+                                 std::array<mult_func_t, num_kernels>,
+                                 std::index_sequence<idx...> const&)
+      {
+         (create_action_callback(mfem::get<idx>(ks), funcs[idx]), ...);
+      }
+
+      Action(DifferentiableOperator &op, kernels_tuple &ks) : op(op)
+      {
+         materialize_callbacks(ks, funcs,
+                               std::make_index_sequence<mfem::tuple_size<kernels_tuple>::value>());
+      }
+
+      void Mult(const Vector &x, Vector &y) const
+      {
+         prolongation(op.solutions, x, solutions_l);
+
+         residual_e = 0.0;
+         for (const auto &f : funcs)
+         {
+            f(residual_e);
+         }
+
+         prolongation_transpose(residual_l, y);
+
+         y.SetSubVector(op.ess_tdof_list, 0.0);
+      }
+
+      void SetParameters(std::vector<Vector *> p) const
+      {
+         MFEM_ASSERT(num_parameters == p.size(),
+                     "number of parameters doesn't match descriptors");
+         for (int i = 0; i < num_parameters; i++)
+         {
+            p[i]->Read();
+            parameters_l[i] = *p[i];
+            // parameters_l[i].MakeRef(p[i], 0, p[i]->Size());
+         }
+      }
+
+   protected:
+      DifferentiableOperator &op;
+      std::array<mult_func_t, num_kernels> funcs;
+
+      std::function<void(Vector &, Vector &)> prolongation_transpose;
+
+      mutable std::array<Vector, num_solutions> solutions_l;
+      mutable std::array<Vector, num_parameters> parameters_l;
+      mutable Vector residual_l;
+
+      mutable std::array<Vector, num_fields> fields_e;
+      mutable Vector residual_e;
+   };
+
+   template <size_t derivative_idx>
+   class Derivative : public Operator
+   {
+   public:
+      template <typename kernel_t>
+      void create_callback(kernel_t kernel, mult_func_t &func);
+
+      template<std::size_t... idx>
+      void materialize_callbacks(kernels_tuple &ks,
+                                 std::array<mult_func_t, num_kernels>,
+                                 std::index_sequence<idx...> const&)
+      {
+         (create_callback(mfem::get<idx>(ks), funcs[idx]), ...);
+      }
+
+      Derivative(
+         DifferentiableOperator &op,
+         std::array<Vector *, num_solutions> &solutions,
+         std::array<Vector *, num_parameters> &parameters,
+         kernels_tuple &ks) : op(op), ks(ks)
+      {
+         for (int i = 0; i < num_solutions; i++)
+         {
+            solutions_l[i] = *solutions[i];
+         }
+
+         for (int i = 0; i < num_parameters; i++)
+         {
+            parameters_l[i] = *parameters[i];
+         }
+
+         // G
+         // if constexpr (std::is_same_v<OperatesOn, OperatesOnElement>)
+         // {
+         element_restriction(op.solutions, solutions_l, fields_e,
+                             op.element_dof_ordering);
+         element_restriction(op.parameters, parameters_l, fields_e,
+                             op.element_dof_ordering,
+                             op.solutions.size());
+         // }
+         // else
+         // {
+         //    MFEM_ABORT("restriction not implemented for OperatesOn");
+         // }
+         direction = op.fields[derivative_idx];
+
+         size_t derivative_action_l_size = 0;
+         for (auto &s : op.solutions)
+         {
+            derivative_action_l_size += GetVSize(s);
+            this->width += GetTrueVSize(s);
+         }
+         this->height = derivative_action_l_size;
+         derivative_action_l.SetSize(derivative_action_l_size);
+
+         materialize_callbacks(ks, funcs,
+                               std::make_index_sequence<num_kernels>());
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         current_direction_t = x;
+         current_direction_t.SetSubVector(op.ess_tdof_list, 0.0);
+
+         prolongation(direction, current_direction_t, direction_l);
+
+         derivative_action_e = 0.0;
+         for (const auto &f : funcs)
+         {
+            f(derivative_action_e);
+         }
+
+         prolongation_transpose(derivative_action_l, y);
+
+         y.SetSubVector(op.ess_tdof_list, 0.0);
+      }
+
+      template <typename kernel_t>
+      void assemble_vector_impl(kernel_t kernel, Vector &v);
+
+      template<std::size_t... idx>
+      void assemble_vector(
+         kernels_tuple &ks,
+         Vector &v,
+         std::index_sequence<idx...> const&)
+      {
+         (assemble_vector_impl(mfem::get<idx>(ks), v), ...);
+      }
+
+      void Assemble(Vector &v)
+      {
+         assemble_vector(ks, v, std::make_index_sequence<num_kernels>());
+      }
+
+      template <typename kernel_t>
+      void assemble_hypreparmatrix_impl(kernel_t kernel, HypreParMatrix &A);
+
+      template<std::size_t... idx>
+      void assemble_hypreparmatrix(
+         kernels_tuple &ks,
+         HypreParMatrix &A,
+         std::index_sequence<idx...> const&)
+      {
+         (assemble_hypreparmatrix_impl(mfem::get<idx>(ks), A), ...);
+      }
+
+      void Assemble(HypreParMatrix &A)
+      {
+         assemble_hypreparmatrix(ks, A, std::make_index_sequence<num_kernels>());
+      }
+
+      void AssembleDiagonal(Vector &d) const override {}
+
+   protected:
+      DifferentiableOperator &op;
+      kernels_tuple &ks;
+      std::array<mult_func_t, num_kernels> funcs;
+
+      std::function<void(Vector &, Vector &)> prolongation_transpose;
+
+      FieldDescriptor direction;
+
+      std::array<Vector, num_solutions> solutions_l;
+      std::array<Vector, num_parameters> parameters_l;
+      mutable Vector direction_l;
+      mutable Vector derivative_action_l;
+
+      mutable std::array<Vector, num_fields> fields_e;
+      mutable Vector direction_e;
+      mutable Vector derivative_action_e;
+
+      mutable Vector current_direction_t;
+   };
+
+   DifferentiableOperator(std::array<FieldDescriptor, num_solutions> s,
+                          std::array<FieldDescriptor, num_parameters> p,
+                          kernels_tuple ks,
+                          ParMesh &m,
+                          autodiff_t ad = AutoDiff::NativeDualNumber{}) :
+      kernels(ks),
+      mesh(m),
+      dim(mesh.Dimension()),
+      solutions(s),
+      parameters(p)
+   {
+      for (int i = 0; i < num_solutions; i++)
+      {
+         fields[i] = solutions[i];
+      }
+
+      for (int i = 0; i < num_parameters; i++)
+      {
+         fields[i + num_solutions] = parameters[i];
+      }
+
+      residual.reset(new Action(*this, kernels));
+   }
+
+   void SetParameters(std::vector<Vector *> p) const
+   {
+      residual->SetParameters(p);
+   }
+
+   void Mult(const Vector &x, Vector &y) const override
+   {
+      residual->Mult(x, y);
+   }
+
+   template <int derivative_idx>
+   std::shared_ptr<Derivative<derivative_idx>>
+                                            GetDerivativeWrt(std::array<Vector *, num_solutions> solutions,
+                                                             std::array<Vector *, num_parameters> parameters)
+   {
+      return std::shared_ptr<Derivative<derivative_idx>>(
+                new Derivative<derivative_idx>(*this, solutions, parameters, kernels));
+   }
+
+   void SetEssentialTrueDofs(const Array<int> &l)
+   {
+      l.Copy(ess_tdof_list);
+   }
+
+   kernels_tuple kernels;
+   ParMesh &mesh;
+   const int dim;
+
+   std::array<FieldDescriptor, num_solutions> solutions;
+   std::array<FieldDescriptor, num_parameters> parameters;
+   // solutions and parameters
+   std::array<FieldDescriptor, num_fields> fields;
+
+   int residual_lsize = 0;
+
+   mutable std::array<Vector, num_solutions> current_state_l;
+   mutable Vector direction_l;
+
+   mutable Vector current_direction_t;
+
+   Array<int> ess_tdof_list;
+
+   static constexpr ElementDofOrdering element_dof_ordering =
+      ElementDofOrdering::LEXICOGRAPHIC;
+
+   static constexpr DofToQuad::Mode doftoquad_mode =
+      DofToQuad::Mode::TENSOR;
+
+   // static constexpr ElementDofOrdering element_dof_ordering =
+   //    ElementDofOrdering::NATIVE;
+
+   // static constexpr DofToQuad::Mode doftoquad_mode =
+   //    DofToQuad::Mode::FULL;
+
+   std::shared_ptr<Action> residual;
+};
+
+template <
+   typename kernels_tuple,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t num_fields,
+   size_t num_kernels,
+   typename autodiff_t
+   >
+template <
+   typename kernel_t
+   >
+void DifferentiableOperator<kernels_tuple,
+     num_solutions,
+     num_parameters,
+     num_fields,
+     num_kernels,
+     autodiff_t>::Action::create_action_callback(
+        kernel_t kernel,
+        mult_func_t &func)
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                     kernel.inputs, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                      kernel.outputs, std::make_index_sequence<kernel.num_koutputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(op.fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(kernel.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
+   const int num_entities = GetNumEntities<entity_t>(op.mesh);
+   const int num_qp = kernel.integration_rule.GetNPoints();
+
+   // All solutions T-vector sizes make up the width of the operator, since
+   // they are explicitly provided in Mult() for example.
+
+   op.width = GetTrueVSize(op.fields[test_space_field_idx]);
+   op.residual_lsize = GetVSize(op.fields[test_space_field_idx]);
+
+   if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      op.height = 1;
+   }
+   else
+   {
+      op.height = op.residual_lsize;
+   }
+
+   residual_l.SetSize(op.residual_lsize);
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : op.fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, kernel.integration_rule,
+                                              doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/op.mesh.Dimension()) + 0.5);
+
+   residual_e.SetSize(R->Height());
+
+   const int residual_size_on_qp = GetSizeOnQP<entity_t>(
+                                      mfem::get<hardcoded_output_idx>(kernel.outputs),
+                                      op.fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(kernel.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(kernel.outputs, dtq,
+                                                    koutput_to_field);
+
+   auto input_fops = create_bare_fops(kernel.inputs);
+   auto output_fops = create_bare_fops(kernel.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(output_fops).size_on_qp /
+      mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(output_fops).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(kernel.integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
+                                                std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto shmem_info = get_shmem_info<entity_t>(input_dtq_maps,
+                                              output_dtq_maps,
+                                              op.fields,
+                                              num_entities,
+                                              kernel.inputs,
+                                              num_qp,
+                                              input_size_on_qp,
+                                              residual_size_on_qp);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   // print_shared_memory_info(shmem_info);
+
+   func = [=](Vector &ye_mem) mutable
+   {
+      restriction<entity_t>(op.solutions, solutions_l, this->fields_e,
+                            op.element_dof_ordering);
+      restriction<entity_t>(op.parameters, parameters_l, this->fields_e,
+                            op.element_dof_ordering,
+                            op.solutions.size());
+
+      auto ye = Reshape(ye_mem.ReadWrite(), test_vdim, num_test_dof, num_entities);
+      auto wrapped_fields_e = wrap_fields(this->fields_e, shmem_info.field_sizes, num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, void *shmem)
+      {
+         // printf("\ne: %d\n", e);
+         // tic();
+         auto input_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT_DTQ],
+            shmem_info.input_dtq_sizes,
+            input_dtq_maps);
+
+         auto output_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT_DTQ],
+            shmem_info.output_dtq_sizes,
+            output_dtq_maps);
+
+         auto fields_shmem = load_field_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::FIELD],
+            shmem_info.field_sizes,
+            kinput_to_field,
+            input_fops,
+            wrapped_fields_e,
+            e,
+         std::make_index_sequence<kernel.num_kinputs> {});
+
+         // These functions don't copy, they simply create a `DeviceTensor` object
+         // that points to correct chunks of the shared memory pool.
+         auto input_shmem = load_input_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::INPUT],
+                               shmem_info.input_sizes,
+                               num_qp);
+
+         auto residual_shmem = load_residual_mem(
+                                  shmem,
+                                  shmem_info.offsets[SharedMemory::Index::OUTPUT],
+                                  shmem_info.residual_size,
+                                  num_qp);
+
+         auto scratch_mem = load_scratch_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::TEMP],
+                               shmem_info.temp_sizes);
+
+         MFEM_SYNC_THREAD;
+         // printf("shmem load elapsed: %.1fus\n", toc() * 1e6);
+
+         // tic();
+         map_fields_to_quadrature_data<TensorProduct>(
+            input_shmem, fields_shmem, input_dtq_shmem, input_fops, ir_weights, scratch_mem,
+            std::make_index_sequence<kernel.num_kinputs> {});
+         // printf("interpolate elapsed: %.1fus\n", toc() * 1e6);
+
+         // tic();
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  const int q = qx + q1d * (qy + q1d * qz);
+                  auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+                  auto r = Reshape(&residual_shmem(0, q), residual_size_on_qp);
+                  apply_kernel(r, kernel.func, kernel_args, input_shmem, q);
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+         // printf("qf elapsed: %.1fus\n", toc() * 1e6);
+
+         // tic();
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                      mfem::get<0>(output_fops),
+                                                      output_dtq_shmem[hardcoded_output_idx],
+                                                      scratch_mem);
+         // printf("integrate elapsed: %.1fus\n", toc() * 1e6);
+
+      }, num_entities, q1d, q1d, q1d, shmem_info.total_size, shmem_cache.ReadWrite());
+
+      if constexpr (std::is_same_v<decltype(output_fop), None>)
+      {
+         residual_l = ye_mem;
+      }
+      else
+      {
+         R->MultTranspose(ye_mem, residual_l);
+      }
+   };
+
+   if constexpr (std::is_same_v<decltype(output_fop), None>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         y = r_local;
+      };
+   }
+   else if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         double local_sum = r_local.Sum();
+         MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM,
+                       op.mesh.GetComm());
+         MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      };
+   }
+   else
+   {
+      auto P = get_prolongation(op.fields[test_space_field_idx]);
+      prolongation_transpose = [P](const Vector &r_local, Vector &y)
+      {
+         P->MultTranspose(r_local, y);
+      };
+   }
+}
+
+template <
+   typename kernels_tuple,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t num_fields,
+   size_t num_kernels,
+   typename autodiff_t
+   >
+template <
+   size_t derivative_idx
+   >
+template <
+   typename kernel_t
+   >
+void DifferentiableOperator<kernels_tuple,
+     num_solutions,
+     num_parameters,
+     num_fields,
+     num_kernels,
+     autodiff_t>::Derivative<derivative_idx>::create_callback(kernel_t kernel,
+                                                              mult_func_t &func)
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                     kernel.inputs, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+                                                                      kernel.outputs, std::make_index_sequence<kernel.num_koutputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(op.fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(kernel.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
+   const int num_entities = GetNumEntities<entity_t>(op.mesh);
+   const int num_qp = kernel.integration_rule.GetNPoints();
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : op.fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, kernel.integration_rule,
+                                              doftoquad_mode));
+   }
+   const int q1d = dtq[0]->nqpt;
+
+   derivative_action_e.SetSize(R->Height());
+
+   const int da_size_on_qp = GetSizeOnQP<entity_t>(
+                                mfem::get<hardcoded_output_idx>(kernel.outputs),
+                                op.fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(kernel.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(kernel.outputs, dtq,
+                                                    koutput_to_field);
+
+   auto input_fops = create_bare_fops(kernel.inputs);
+   auto output_fops = create_bare_fops(kernel.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(output_fops).size_on_qp /
+      mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(output_fops).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(kernel.integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
+                                                std::make_index_sequence<kernel.num_kinputs> {});
+
+   // Check which qf inputs are dependent on the dependent variable
+   std::array<bool, kernel.num_kinputs> kinput_is_dependent;
+   bool no_kinput_is_dependent = true;
+   for (int i = 0; i < kinput_is_dependent.size(); i++)
+   {
+      if (kinput_to_field[i] == derivative_idx)
+      {
+         no_kinput_is_dependent = false;
+         kinput_is_dependent[i] = true;
+         // out << "function input " << i << " is dependent on "
+         //     << op.fields[kinput_to_field[i]].field_label << "\n";
+      }
+      else
+      {
+         kinput_is_dependent[i] = false;
+      }
+   }
+
+   bool with_derivatives = true;
+   auto shmem_info = get_shmem_info<entity_t>(input_dtq_maps,
+                                              output_dtq_maps,
+                                              op.fields,
+                                              num_entities,
+                                              kernel.inputs,
+                                              num_qp,
+                                              input_size_on_qp,
+                                              da_size_on_qp,
+                                              derivative_idx);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   // print_shared_memory_info(shmem_info);
+
+   func = [=](Vector &ye_mem) mutable
+   {
+      if (no_kinput_is_dependent)
+      {
+         return;
+      }
+
+      restriction<entity_t>(direction, direction_l, direction_e,
+                            op.element_dof_ordering);
+
+      auto ye = Reshape(ye_mem.ReadWrite(), num_test_dof, test_vdim, num_entities);
+      auto wrapped_fields_e = wrap_fields(this->fields_e, shmem_info.field_sizes, num_entities);
+      auto wrapped_direction_e = Reshape(direction_e.ReadWrite(), shmem_info.direction_size, num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, double *shmem)
+      {
+         auto input_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT_DTQ],
+            shmem_info.input_dtq_sizes,
+            input_dtq_maps);
+
+         auto output_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT_DTQ],
+            shmem_info.output_dtq_sizes,
+            output_dtq_maps);
+
+         auto fields_shmem = load_field_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::FIELD],
+            shmem_info.field_sizes,
+            kinput_to_field,
+            input_fops,
+            wrapped_fields_e,
+            e,
+         std::make_index_sequence<kernel.num_kinputs> {});
+
+         auto direction_shmem = load_direction_mem(
+                                   shmem,
+                                   shmem_info.offsets[SharedMemory::Index::DIRECTION],
+                                   shmem_info.direction_size,
+                                   wrapped_direction_e,
+                                   e);
+
+         // These methods don't copy, they simply create a `DeviceTensor` object
+         // that points to correct chunks of the shared memory pool.
+         auto input_shmem = load_input_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::INPUT],
+                               shmem_info.input_sizes,
+                               num_qp);
+
+         auto shadow_shmem = load_input_mem(
+                                shmem,
+                                shmem_info.offsets[SharedMemory::Index::SHADOW],
+                                shmem_info.input_sizes,
+                                num_qp);
+
+         auto residual_shmem = load_residual_mem(
+                                  shmem,
+                                  shmem_info.offsets[SharedMemory::Index::OUTPUT],
+                                  shmem_info.residual_size,
+                                  num_qp);
+
+         auto scratch_mem = load_scratch_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::TEMP],
+                               shmem_info.temp_sizes);
+
+         map_fields_to_quadrature_data<TensorProduct>(
+            input_shmem, fields_shmem, input_dtq_shmem, input_fops, ir_weights, scratch_mem,
+            std::make_index_sequence<kernel.num_kinputs> {});
+
+         zero_all(shadow_shmem);
+         map_direction_to_quadrature_data_conditional<TensorProduct>(
+            shadow_shmem, direction_shmem, input_dtq_shmem, input_fops, ir_weights,
+            scratch_mem, kinput_is_dependent,
+            std::make_index_sequence<kernel.num_kinputs> {});
+
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  const int q = qx + q1d * (qy + q1d * qz);
+                  auto r = Reshape(&residual_shmem(0, q), da_size_on_qp);
+                  auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+
+                  if constexpr (std::is_same_v<autodiff_t, AutoDiff::EnzymeForward>)
+                  {
+                     auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+                     apply_kernel_fwddiff_enzyme(
+                        r,
+                        kernel.func,
+                        kernel_args,
+                        kernel_shadow_args,
+                        input_shmem,
+                        shadow_shmem,
+                        q);
+                  }
+                  else if constexpr (std::is_same_v<autodiff_t, AutoDiff::NativeDualNumber>)
+                  {
+                     apply_kernel_native_dual(
+                        r,
+                        kernel.func,
+                        kernel_args,
+                        input_shmem,
+                        shadow_shmem,
+                        q);
+                  }
+                  else
+                  {
+                     static_assert(always_false<autodiff_t>, "unknown autodiff type");
+                  }
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                      mfem::get<0>(output_fops),
+                                                      output_dtq_shmem[hardcoded_output_idx],
+                                                      scratch_mem);
+      }, num_entities, q1d, q1d, 1, shmem_info.total_size, shmem_cache.ReadWrite());
+
+      R->MultTranspose(ye_mem, derivative_action_l);
+   };
+
+   if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         double local_sum = r_local.Sum();
+         MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM,
+                       op.mesh.GetComm());
+         MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      };
+   }
+   else
+   {
+      auto P = get_prolongation(op.fields[test_space_field_idx]);
+      prolongation_transpose = [P](const Vector &r_local, Vector &y)
+      {
+         P->MultTranspose(r_local, y);
+      };
+   }
+}
+
+} // namespace mfem

examples/dfem/dfem_element_operator.hpp

@@ -0,0 +1,79 @@
+#include "dfem_util.hpp"
+
+namespace mfem
+{
+
+template <typename func_t, typename input_t, typename output_t, typename dependency_map_t>
+struct ElementOperator;
+
+template <typename func_t, typename... input_ts, typename... output_ts, typename dependency_map_t>
+struct ElementOperator<func_t, mfem::tuple<input_ts...>, mfem::tuple<output_ts...>, dependency_map_t>
+{
+   using entity_t = Entity::Element;
+
+   func_t qfunc;
+
+   mfem::tuple<input_ts...> inputs;
+   mfem::tuple<output_ts...> outputs;
+
+   dependency_map_t dependency_map;
+
+   using qf_param_ts = typename create_function_signature<
+                       decltype(&func_t::operator())>::type::parameter_ts;
+   using qf_output_t = typename create_function_signature<
+                       decltype(&func_t::operator())>::type::return_t;
+
+   static constexpr size_t num_inputs =
+      mfem::tuple_size<decltype(inputs)>::value;
+   static constexpr size_t num_outputs =
+      mfem::tuple_size<decltype(outputs)>::value;
+
+   ElementOperator(func_t qfunc,
+                   mfem::tuple<input_ts...> inputs,
+                   mfem::tuple<output_ts...> outputs)
+      : qfunc(qfunc), inputs(inputs), outputs(outputs),
+        dependency_map(make_dependency_map(inputs))
+   {
+      // Consistency checks
+      if constexpr (num_outputs > 1)
+      {
+         static_assert(always_false<func_t>,
+                       "more than one output per kernel is not supported right now");
+      }
+
+      constexpr size_t num_qfinputs = mfem::tuple_size<qf_param_ts>::value;
+      static_assert(num_qfinputs == num_inputs,
+                    "kernel function inputs and descriptor inputs have to match");
+
+      constexpr size_t num_qf_outputs = mfem::tuple_size<qf_output_t>::value;
+      static_assert(num_qf_outputs == num_qf_outputs,
+                    "kernel function outputs and descriptor outputs have to match");
+   }
+};
+
+template <typename func_t, typename... input_ts, typename... output_ts>
+ElementOperator(func_t, mfem::tuple<input_ts...>, mfem::tuple<output_ts...>)
+-> ElementOperator<func_t, mfem::tuple<input_ts...>, mfem::tuple<output_ts...>,
+decltype(make_dependency_map(std::declval<mfem::tuple<input_ts...>>()))>;
+
+// template <typename func_t, typename input_t, typename output_t>
+// struct BoundaryElementOperator : public
+//    ElementOperator<func_t, input_t, output_t>
+// {
+// public:
+//    using entity_t = Entity::BoundaryElement;
+//    BoundaryElementOperator(func_t func, input_t inputs, output_t outputs)
+//       : ElementOperator<func_t, input_t, output_t>(func, inputs, outputs) {}
+// };
+
+// template <typename func_t, typename input_t, typename output_t>
+// struct FaceOperator : public
+//    ElementOperator<func_t, input_t, output_t>
+// {
+// public:
+//    using entity_t = Entity::Face;
+//    FaceOperator(func_t func, input_t inputs, output_t outputs)
+//       : ElementOperator<func_t, input_t, output_t>(func, inputs, outputs) {}
+// };
+
+} // namespace mfem

examples/dfem/dfem_fieldoperator.hpp

@@ -0,0 +1,246 @@
+#pragma once
+
+#include <string>
+
+namespace mfem
+{
+
+template <int FIELD_ID = -1>
+class FieldOperator
+{
+public:
+   constexpr FieldOperator(int size_on_qp = 0) :
+      size_on_qp(size_on_qp) {};
+
+   static constexpr int GetFieldId() { return FIELD_ID; }
+
+   int size_on_qp = -1;
+
+   int dim = -1;
+
+   int vdim = -1;
+};
+
+template <int FIELD_ID = -1>
+class None : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr None() : FieldOperator<FIELD_ID>() {}
+};
+
+template< typename T >
+struct is_none_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_none_fop<None<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <typename T>
+struct DisableAD
+{
+   T& operator()() const { return fop; }
+   T fop;
+};
+
+class Weight : public FieldOperator<-1>
+{
+public:
+   constexpr Weight() : FieldOperator<-1>() {};
+};
+
+template< typename T >
+struct is_weight_fop
+{
+   static const bool value = false;
+};
+
+template <>
+struct is_weight_fop<Weight>
+{
+   static const bool value = true;
+};
+
+template <int FIELD_ID = -1>
+class Value : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr Value() : FieldOperator<FIELD_ID>() {};
+};
+
+template< typename T >
+struct is_value_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_value_fop<Value<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <typename T>
+struct is_value_fop<DisableAD<T>>
+{
+   static const bool value = is_value_fop<T>::value;
+};
+
+template <int FIELD_ID = -1>
+class Gradient : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr Gradient() : FieldOperator<FIELD_ID>() {};
+};
+
+template< typename T >
+struct is_gradient_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_gradient_fop<Gradient<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+template <int FIELD_ID = -1>
+class One : public FieldOperator<FIELD_ID>
+{
+public:
+   constexpr One() : FieldOperator<FIELD_ID>() {};
+};
+
+template< typename T >
+struct is_one_fop
+{
+   static const bool value = false;
+};
+
+template <int FIELD_ID>
+struct is_one_fop<One<FIELD_ID>>
+{
+   static const bool value = true;
+};
+
+// class FieldOperator
+// {
+// public:
+//    FieldOperator(std::string field_label = "", int size_on_qp = 0) :
+//       field_label(field_label),
+//       size_on_qp(size_on_qp) {};
+
+//    std::string field_label;
+
+//    int size_on_qp = -1;
+
+//    int dim = -1;
+
+//    int vdim = -1;
+// };
+
+// class None : public FieldOperator
+// {
+// public:
+//    None(std::string field_label) :
+//       FieldOperator(field_label) {}
+// };
+
+// class Weight : public FieldOperator
+// {
+// public:
+//    Weight() : FieldOperator("quadrature_weights") {};
+// };
+
+// class Value : public FieldOperator
+// {
+// public:
+//    Value(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class Gradient : public FieldOperator
+// {
+// public:
+//    Gradient(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class Curl : public FieldOperator
+// {
+// public:
+//    Curl(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class Div : public FieldOperator
+// {
+// public:
+//    Div(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class FaceValueLeft : public FieldOperator
+// {
+// public:
+//    FaceValueLeft(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class FaceValueRight : public FieldOperator
+// {
+// public:
+//    FaceValueRight(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class FaceNormal : public FieldOperator
+// {
+// public:
+//    FaceNormal(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// class One : public FieldOperator
+// {
+// public:
+//    One(std::string field_label) : FieldOperator(field_label) {};
+// };
+
+// namespace BareFieldOperator
+// {
+
+// struct Base
+// {
+//    Base(FieldOperator &o)
+//    {
+//       size_on_qp = o.size_on_qp;
+//       dim = o.dim;
+//       vdim = o.vdim;
+//    };
+//    int size_on_qp = -1;
+//    int dim = -1;
+//    int vdim = -1;
+// };
+
+// struct None : Base
+// {
+//    None(FieldOperator &o) : Base(o) {}
+// };
+
+// struct Weight : Base
+// {
+//    Weight(FieldOperator &o) : Base(o) {}
+// };
+
+// struct Value : Base
+// {
+//    Value(FieldOperator &o) : Base(o) {}
+// };
+
+// struct Gradient : Base
+// {
+//    Gradient(FieldOperator &o) : Base(o) {}
+// };
+
+// }
+
+} // namespace mfem

examples/dfem/dfem_integrate.hpp

@@ -0,0 +1,426 @@
+#pragma once
+
+#include "dfem_util.hpp"
+#include <type_traits>
+
+namespace mfem
+{
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields_impl(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+   // assuming the quadrature point residual has to "play nice with
+   // the test function"
+   if constexpr (is_value_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [num_qp, cdim, num_dof] = B.GetShape();
+      const int vdim = output.vdim > 0 ? output.vdim : cdim ;
+      for (int dof = 0; dof < num_dof; dof++)
+      {
+         for (int vd = 0; vd < vdim; vd++)
+         {
+            double acc = 0.0;
+            for (int qp = 0; qp < num_qp; qp++)
+            {
+               acc += B(qp, 0, dof) * f(vd, 0, qp);
+            }
+            y(dof, vd) += acc;
+         }
+      }
+   }
+   else if constexpr (
+      is_gradient_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [num_qp, dim, num_dof] = G.GetShape();
+      const int vdim = output.vdim;
+      for (int dof = 0; dof < num_dof; dof++)
+      {
+         for (int vd = 0; vd < vdim; vd++)
+         {
+            double acc = 0.0;
+            for (int d = 0; d < dim; d++)
+            {
+               for (int qp = 0; qp < num_qp; qp++)
+               {
+                  acc += G(qp, d, dof) * f(vd, d, qp);
+               }
+            }
+            y(dof, vd) += acc;
+         }
+      }
+   }
+   else if constexpr (is_one_fop<std::decay_t<output_t>>::value)
+   {
+      // This is the "integral over all quadrature points type" applying
+      // B = 1 s.t. B^T * C \in R^1.
+      const auto [num_qp, unused, unused1] = B.GetShape();
+      auto cc = Reshape(&f(0, 0, 0), num_qp);
+      for (int i = 0; i < num_qp; i++)
+      {
+         y(0, 0) += cc(i);
+      }
+   }
+   else if constexpr (is_none_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [num_qp, unused, num_dof] = B.GetShape();
+      const auto vdim = output.vdim;
+      auto cc = Reshape(&f(0, 0, 0), num_qp * vdim);
+      auto yy = Reshape(&y(0, 0), num_qp * vdim);
+      for (int i = 0; i < num_qp * vdim; i++)
+      {
+         yy(i) = cc(i);
+      }
+   }
+   else
+   {
+      MFEM_ABORT("quadrature data mapping to field is not implemented for"
+                 " this field descriptor");
+   }
+}
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields_tensor_impl_2d(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq,
+   std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               double acc = 0.0;
+               for (int qx = 0; qx < q1d; qx++)
+               {
+                  acc += fqp(vd, 0, qx, qy) * B(qx, 0, dx);
+               }
+               s0(qy, dx) = acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               double acc = 0.0;
+               for (int qy = 0; qy < q1d; qy++)
+               {
+                  acc += s0(qy, dx) * B(qy, 0, dy);
+               }
+               yd(dx, dy, vd) += acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_gradient_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = G.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, d1d);
+      auto s1 = Reshape(&scratch_mem[1](0), q1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int qx = 0; qx < q1d; qx++)
+               {
+                  uv[0] += fqp(vd, 0, qx, qy) * G(qx, 0, dx);
+                  uv[1] += fqp(vd, 1, qx, qy) * B(qx, 0, dx);
+               }
+               s0(qy, dx) = uv[0];
+               s1(qy, dx) = uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int qy = 0; qy < q1d; qy++)
+               {
+                  uv[0] += s0(qy, dx) * B(qy, 0, dy);
+                  uv[1] += s1(qy, dx) * G(qy, 0, dy);
+               }
+               yd(dx, dy, vd) += uv[0] + uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_none_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      auto fqp = Reshape(&f(0, 0, 0), output.size_on_qp, q1d, q1d);
+      auto yqp = Reshape(&y(0, 0), output.size_on_qp, q1d, q1d);
+
+      for (int sq = 0; sq < output.size_on_qp; sq++)
+      {
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               yqp(sq, qx, qy) = fqp(sq, qx, qy);
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else
+   {
+      MFEM_ABORT("quadrature data mapping to field is not implemented for"
+                 " this field descriptor with sum factorization on tensor product elements");
+   }
+}
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields_tensor_impl_3d(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq,
+   std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, q1d, d1d);
+      auto s1 = Reshape(&scratch_mem[1](0), q1d, d1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  double acc = 0.0;
+                  for (int qx = 0; qx < q1d; qx++)
+                  {
+                     acc += fqp(vd, 0, qx, qy, qz) * B(qx, 0, dx);
+                  }
+                  s0(qz, qy, dx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  double acc = 0.0;
+                  for (int qy = 0; qy < q1d; qy++)
+                  {
+                     acc += s0(qz, qy, dx) * B(qy, 0, dy);
+                  }
+                  s1(qz, dy, dx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(dx, x, d1d)
+            {
+               MFEM_FOREACH_THREAD(dz, z, d1d)
+               {
+                  double acc = 0.0;
+                  for (int qz = 0; qz < q1d; qz++)
+                  {
+                     acc += s1(qz, dy, dx) * B(qz, 0, dz);
+                  }
+                  yd(dx, dy, dz, vd) += acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_gradient_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = G.GetShape();
+      const int vdim = output.vdim;
+      const int test_dim = output.size_on_qp / vdim;
+      auto fqp = Reshape(&f(0, 0, 0), vdim, test_dim, q1d, q1d, q1d);
+      auto yd = Reshape(&y(0, 0), d1d, d1d, d1d, vdim);
+
+      auto s0 = Reshape(&scratch_mem[0](0), q1d, q1d, d1d);
+      auto s1 = Reshape(&scratch_mem[1](0), q1d, q1d, d1d);
+      auto s2 = Reshape(&scratch_mem[2](0), q1d, q1d, d1d);
+      auto s3 = Reshape(&scratch_mem[3](0), q1d, d1d, d1d);
+      auto s4 = Reshape(&scratch_mem[4](0), q1d, d1d, d1d);
+      auto s5 = Reshape(&scratch_mem[5](0), q1d, d1d, d1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(dx, x, d1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int qx = 0; qx < q1d; qx++)
+                  {
+                     uvw[0] += fqp(vd, 0, qx, qy, qz) * G(qx, 0, dx);
+                     uvw[1] += fqp(vd, 1, qx, qy, qz) * B(qx, 0, dx);
+                     uvw[2] += fqp(vd, 2, qx, qy, qz) * B(qx, 0, dx);
+                  }
+                  s0(qz, qy, dx) = uvw[0];
+                  s1(qz, qy, dx) = uvw[1];
+                  s2(qz, qy, dx) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(dx, x, d1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int qy = 0; qy < q1d; qy++)
+                  {
+                     uvw[0] += s0(qz, qy, dx) * B(qy, 0, dy);
+                     uvw[1] += s1(qz, qy, dx) * G(qy, 0, dy);
+                     uvw[2] += s2(qz, qy, dx) * B(qy, 0, dy);
+                  }
+                  s3(qz, dy, dx) = uvw[0];
+                  s4(qz, dy, dx) = uvw[1];
+                  s5(qz, dy, dx) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(dx, x, d1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int qz = 0; qz < q1d; qz++)
+                  {
+                     uvw[0] += s3(qz, dy, dx) * B(qz, 0, dz);
+                     uvw[1] += s4(qz, dy, dx) * B(qz, 0, dz);
+                     uvw[2] += s5(qz, dy, dx) * G(qz, 0, dz);
+                  }
+                  yd(dx, dy, dz, vd) += uvw[0] + uvw[1] + uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (is_none_fop<std::decay_t<output_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      auto fqp = Reshape(&f(0, 0, 0), output.size_on_qp, q1d, q1d, q1d);
+      auto yqp = Reshape(&y(0, 0), output.size_on_qp, q1d, q1d, q1d);
+
+      for (int sq = 0; sq < output.size_on_qp; sq++)
+      {
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  yqp(sq, qx, qy, qz) = fqp(sq, qx, qy, qz);
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else
+   {
+      MFEM_ABORT("quadrature data mapping to field is not implemented for"
+                 " this field descriptor with sum factorization on tensor product elements");
+   }
+}
+
+template <typename output_t>
+MFEM_HOST_DEVICE
+void map_quadrature_data_to_fields(
+   DeviceTensor<2, double> &y,
+   const DeviceTensor<3, double> &f,
+   const output_t &output,
+   const DofToQuadMap &dtq,
+   std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const int &dimension,
+   const bool &use_sum_factorization)
+{
+   if (use_sum_factorization)
+   {
+      if (dimension == 2)
+      {
+         map_quadrature_data_to_fields_tensor_impl_2d(y, f, output, dtq, scratch_mem);
+      }
+      else if (dimension == 3)
+      {
+         map_quadrature_data_to_fields_tensor_impl_3d(y, f, output, dtq, scratch_mem);
+      }
+   }
+   else
+   {
+      map_quadrature_data_to_fields_impl(y, f, output, dtq);
+   }
+}
+
+}

examples/dfem/dfem_interpolate.hpp

@@ -0,0 +1,564 @@
+#pragma once
+
+#include "dfem_util.hpp"
+
+namespace mfem
+{
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE inline
+void map_field_to_quadrature_data_tensor_product_3d(
+   DeviceTensor<2> &field_qp,
+   const DofToQuadMap &dtq,
+   const DeviceTensor<1> &field_e,
+   const field_operator_t &input,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<field_operator_t>>::value)
+   {
+      auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, q1d, q1d, q1d);
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, d1d, q1d);
+      auto s1 = Reshape(&scratch_mem[1](0), d1d, q1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  double acc = 0.0;
+                  for (int dx = 0; dx < d1d; dx++)
+                  {
+                     acc += B(qx, 0, dx) * field(dx, dy, dz, vd);
+                  }
+                  s0(dz, dy, qx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               MFEM_FOREACH_THREAD(qy, y, q1d)
+               {
+                  double acc = 0.0;
+                  for (int dy = 0; dy < d1d; dy++)
+                  {
+                     acc += s0(dz, dy, qx) * B(qy, 0, dy);
+                  }
+                  s1(dz, qy, qx) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  double acc = 0.0;
+                  for (int dz = 0; dz < d1d; dz++)
+                  {
+                     acc += s1(dz, qy, qx) * B(qz, 0, dz);
+                  }
+                  fqp(vd, qx, qy, qz) = acc;
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (
+      is_gradient_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const int dim = input.dim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, dim, q1d, q1d, q1d);
+
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, d1d, q1d);
+      auto s1 = Reshape(&scratch_mem[1](0), d1d, d1d, q1d);
+      auto s2 = Reshape(&scratch_mem[2](0), d1d, q1d, q1d);
+      auto s3 = Reshape(&scratch_mem[3](0), d1d, q1d, q1d);
+      auto s4 = Reshape(&scratch_mem[4](0), d1d, q1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(dy, y, d1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  real_t uv[2] = {0.0, 0.0};
+                  for (int dx = 0; dx < d1d; dx++)
+                  {
+                     const real_t f = field(dx, dy, dz, vd);
+                     uv[0] += f * B(qx, 0, dx);
+                     uv[1] += f * G(qx, 0, dx);
+                  }
+                  s0(dz, dy, qx) = uv[0];
+                  s1(dz, dy, qx) = uv[1];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(dz, z, d1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int dy = 0; dy < d1d; dy++)
+                  {
+                     const real_t s0i = s0(dz, dy, qx);
+                     uvw[0] += s1(dz, dy, qx) * B(qy, 0, dy);
+                     uvw[1] += s0i * G(qy, 0, dy);
+                     uvw[2] += s0i * B(qy, 0, dy);
+                  }
+                  s2(dz, qy, qx) = uvw[0];
+                  s3(dz, qy, qx) = uvw[1];
+                  s4(dz, qy, qx) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qz, z, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qx, x, q1d)
+               {
+                  real_t uvw[3] = {0.0, 0.0, 0.0};
+                  for (int dz = 0; dz < d1d; dz++)
+                  {
+                     uvw[0] += s2(dz, qy, qx) * B(qz, 0, dz);
+                     uvw[1] += s3(dz, qy, qx) * B(qz, 0, dz);
+                     uvw[2] += s4(dz, qy, qx) * G(qz, 0, dz);
+                  }
+                  fqp(vd, 0, qx, qy, qz) = uvw[0];
+                  fqp(vd, 1, qx, qy, qz) = uvw[1];
+                  fqp(vd, 2, qx, qy, qz) = uvw[2];
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   // TODO: Create separate function for clarity
+   else if constexpr (
+      std::is_same_v<std::decay_t<field_operator_t>, Weight>)
+   {
+      const int num_qp = integration_weights.GetShape()[0];
+      // TODO: eeek
+      const int q1d = (int)floor(pow(num_qp, 1.0/input.dim) + 0.5);
+      auto w = Reshape(&integration_weights[0], q1d, q1d, q1d);
+      auto f = Reshape(&field_qp[0], q1d, q1d, q1d);
+      MFEM_FOREACH_THREAD(qx, x, q1d)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(qz, z, q1d)
+            {
+               f(qx, qy, qz) = w(qx, qy, qz);
+            }
+         }
+      }
+      MFEM_SYNC_THREAD;
+   }
+   else if constexpr (is_none_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const int q1d = B.GetShape()[0];
+      auto field = Reshape(&field_e[0], input.size_on_qp, q1d * q1d * q1d);
+      field_qp = field;
+   }
+   else
+   {
+      static_assert(always_false<std::decay_t<field_operator_t>>,
+                    "can't map field to quadrature data");
+   }
+}
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE inline
+void map_field_to_quadrature_data_tensor_product_2d(
+   DeviceTensor<2> &field_qp,
+   const DofToQuadMap &dtq,
+   const DeviceTensor<1> &field_e,
+   const field_operator_t &input,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+
+   if constexpr (is_value_fop<std::decay_t<field_operator_t>>::value)
+   {
+      auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, q1d, q1d);
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               double acc = 0.0;
+               for (int dx = 0; dx < d1d; dx++)
+               {
+                  acc += B(qx, 0, dx) * field(dx, dy, vd);
+               }
+               s0(dy, qx) = acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               double acc = 0.0;
+               for (int dy = 0; dy < d1d; dy++)
+               {
+                  acc += s0(dy, qx) * B(qy, 0, dy);
+               }
+               fqp(vd, qx, qy) = acc;
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   else if constexpr (
+      is_gradient_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const auto [q1d, unused, d1d] = B.GetShape();
+      const int vdim = input.vdim;
+      const int dim = input.dim;
+      const auto field = Reshape(&field_e[0], d1d, d1d, vdim);
+      auto fqp = Reshape(&field_qp[0], vdim, dim, q1d, q1d);
+
+      auto s0 = Reshape(&scratch_mem[0](0), d1d, q1d);
+      auto s1 = Reshape(&scratch_mem[1](0), d1d, q1d);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         MFEM_FOREACH_THREAD(dy, y, d1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int dx = 0; dx < d1d; dx++)
+               {
+                  const real_t f = field(dx, dy, vd);
+                  uv[0] += f * B(qx, 0, dx);
+                  uv[1] += f * G(qx, 0, dx);
+               }
+               s0(dy, qx) = uv[0];
+               s1(dy, qx) = uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            MFEM_FOREACH_THREAD(qx, x, q1d)
+            {
+               real_t uv[2] = {0.0, 0.0};
+               for (int dy = 0; dy < d1d; dy++)
+               {
+                  const real_t s0i = s0(dy, qx);
+                  uv[0] += s1(dy, qx) * B(qy, 0, dy);
+                  uv[1] += s0i * G(qy, 0, dy);
+               }
+               fqp(vd, 0, qx, qy) = uv[0];
+               fqp(vd, 1, qx, qy) = uv[1];
+            }
+         }
+         MFEM_SYNC_THREAD;
+      }
+   }
+   // TODO: Create separate function for clarity
+   else if constexpr (
+      std::is_same_v<std::decay_t<field_operator_t>, Weight>)
+   {
+      const int num_qp = integration_weights.GetShape()[0];
+      // TODO: eeek
+      const int q1d = (int)floor(pow(num_qp, 1.0/input.dim) + 0.5);
+      auto w = Reshape(&integration_weights[0], q1d, q1d);
+      auto f = Reshape(&field_qp[0], q1d, q1d);
+      MFEM_FOREACH_THREAD(qx, x, q1d)
+      {
+         MFEM_FOREACH_THREAD(qy, y, q1d)
+         {
+            f(qx, qy) = w(qx, qy);
+         }
+      }
+      MFEM_SYNC_THREAD;
+   }
+   else if constexpr (is_none_fop<std::decay_t<field_operator_t>>::value)
+   {
+      const int q1d = B.GetShape()[0];
+      auto field = Reshape(&field_e[0], input.size_on_qp, q1d * q1d);
+      field_qp = field;
+   }
+   else
+   {
+      static_assert(always_false<std::decay_t<field_operator_t>>,
+                    "can't map field to quadrature data");
+   }
+}
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE
+void map_field_to_quadrature_data(
+   DeviceTensor<2> field_qp,
+   const DofToQuadMap &dtq,
+   const DeviceTensor<1> &field_e,
+   const field_operator_t &input,
+   const DeviceTensor<1, const double> &integration_weights)
+{
+   auto B = dtq.B;
+   auto G = dtq.G;
+   if constexpr (is_value_fop<field_operator_t>::value)
+   {
+      auto [num_qp, dim, num_dof] = B.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e(0), num_dof, vdim);
+
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         for (int qp = 0; qp < num_qp; qp++)
+         {
+            double acc = 0.0;
+            for (int dof = 0; dof < num_dof; dof++)
+            {
+               acc += B(qp, 0, dof) * field(dof, vd);
+            }
+            field_qp(vd, qp) = acc;
+         }
+      }
+   }
+   else if constexpr (is_gradient_fop<field_operator_t>::value)
+   {
+      const auto [num_qp, dim, num_dof] = G.GetShape();
+      const int vdim = input.vdim;
+      const auto field = Reshape(&field_e(0), num_dof, vdim);
+
+      auto f = Reshape(&field_qp[0], vdim, dim, num_qp);
+      for (int vd = 0; vd < vdim; vd++)
+      {
+         for (int qp = 0; qp < num_qp; qp++)
+         {
+            for (int d = 0; d < dim; d++)
+            {
+               double acc = 0.0;
+               for (int dof = 0; dof < num_dof; dof++)
+               {
+                  acc += G(qp, d, dof) * field(dof, vd);
+               }
+               f(vd, d, qp) = acc;
+            }
+         }
+      }
+   }
+   // else if constexpr (std::is_same_v<field_operator_t, FaceNormal>)
+   // {
+   //    auto normal = geometric_factors.normal;
+   //    auto [num_qp, dim, num_entities] = normal.GetShape();
+   //    auto f = Reshape(&field_qp[0], dim, num_qp);
+   //    for (int qp = 0; qp < num_qp; qp++)
+   //    {
+   //       for (int d = 0; d < dim; d++)
+   //       {
+   //          f(d, qp) = normal(qp, d, entity_idx);
+   //       }
+   //    }
+   // }
+   // TODO: Create separate function for clarity
+   else if constexpr (std::is_same_v<field_operator_t, Weight>)
+   {
+      const int num_qp = integration_weights.GetShape()[0];
+      auto f = Reshape(&field_qp[0], num_qp);
+      for (int qp = 0; qp < num_qp; qp++)
+      {
+         f(qp) = integration_weights(qp);
+      }
+   }
+   else if constexpr (is_none_fop<field_operator_t>::value)
+   {
+      auto [num_qp, unused, num_dof] = B.GetShape();
+      const int size_on_qp = input.size_on_qp;
+      const auto field = Reshape(&field_e[0], size_on_qp * num_qp);
+      auto f = Reshape(&field_qp[0], size_on_qp * num_qp);
+      for (int i = 0; i < size_on_qp * num_qp; i++)
+      {
+         f(i) = field(i);
+      }
+   }
+   else
+   {
+      static_assert(always_false<field_operator_t>,
+                    "can't map field to quadrature data");
+
+   }
+}
+
+template <typename field_operator_ts, size_t num_inputs, size_t num_fields>
+MFEM_HOST_DEVICE inline
+void map_fields_to_quadrature_data(
+   std::array<DeviceTensor<2>, num_inputs> &fields_qp,
+   const std::array<DeviceTensor<1>, num_fields> &fields_e,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
+   const std::array<int, num_inputs> &input_to_field,
+   const field_operator_ts &fops,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const int &dimension,
+   const bool &use_sum_factorization = false)
+{
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      if (use_sum_factorization)
+      {
+         if (dimension == 2)
+         {
+            map_field_to_quadrature_data_tensor_product_2d(
+               fields_qp[i], dtqmaps[i], fields_e[input_to_field[i]], mfem::get<i>(fops),
+               integration_weights, scratch_mem);
+         }
+         else if (dimension == 3)
+         {
+            map_field_to_quadrature_data_tensor_product_3d(
+               fields_qp[i], dtqmaps[i], fields_e[input_to_field[i]], mfem::get<i>(fops),
+               integration_weights, scratch_mem);
+         }
+         else { MFEM_ABORT("unsupported dimension"); }
+      }
+      else
+      {
+         map_field_to_quadrature_data(
+            fields_qp[i], dtqmaps[i], fields_e[input_to_field[i]], mfem::get<i>(fops),
+            integration_weights);
+      }
+   });
+}
+
+template <typename field_operator_t>
+MFEM_HOST_DEVICE
+void map_field_to_quadrature_data_conditional(
+   DeviceTensor<2> &field_qp,
+   const DeviceTensor<1> &field_e,
+   const DofToQuadMap &dtqmap,
+   field_operator_t &fop,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const bool &condition,
+   const int &dimension,
+   const bool &use_sum_factorization = false)
+{
+   if (condition)
+   {
+      if (use_sum_factorization)
+      {
+         if (dimension == 2)
+         {
+            map_field_to_quadrature_data_tensor_product_3d(
+               field_qp, dtqmap, field_e, fop, integration_weights, scratch_mem);
+         }
+         else if (dimension == 3)
+         {
+            map_field_to_quadrature_data_tensor_product_2d(
+               field_qp, dtqmap, field_e, fop, integration_weights, scratch_mem);
+         }
+      }
+      else
+      {
+         map_field_to_quadrature_data(
+            field_qp, dtqmap, field_e, fop, integration_weights);
+      }
+   }
+}
+
+template <size_t num_fields, size_t num_inputs, typename field_operator_ts>
+MFEM_HOST_DEVICE
+void map_fields_to_quadrature_data_conditional(
+   std::array<DeviceTensor<2>, num_inputs> &fields_qp,
+   const std::array<DeviceTensor<1, const double>, num_fields> &fields_e,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
+   field_operator_ts fops,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const std::array<bool, num_inputs> &conditions,
+   const bool &use_sum_factorization = false)
+{
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      map_field_to_quadrature_data_conditional(
+         fields_qp[i], fields_e[i], dtqmaps[i], mfem::get<i>(fops), integration_weights,
+         scratch_mem, conditions[i], use_sum_factorization);
+   });
+}
+
+template <size_t num_inputs, typename field_operator_ts>
+MFEM_HOST_DEVICE
+void map_direction_to_quadrature_data_conditional(
+   std::array<DeviceTensor<2>, num_inputs> &directions_qp,
+   const DeviceTensor<1> &direction_e,
+   const std::array<DofToQuadMap, num_inputs> &dtqmaps,
+   field_operator_ts fops,
+   const DeviceTensor<1, const double> &integration_weights,
+   const std::array<DeviceTensor<1>, 6> &scratch_mem,
+   const std::array<bool, num_inputs> &conditions,
+   const int &dimension,
+   const bool &use_sum_factorization = false)
+{
+   for_constexpr<num_inputs>([&](auto i)
+   {
+      if (conditions[i])
+      {
+         if (use_sum_factorization)
+         {
+            if (dimension == 2)
+            {
+               map_field_to_quadrature_data_tensor_product_2d(
+                  directions_qp[i], dtqmaps[i], direction_e, mfem::get<i>(fops),
+                  integration_weights, scratch_mem);
+            }
+            else if (dimension == 3)
+            {
+               map_field_to_quadrature_data_tensor_product_3d(
+                  directions_qp[i], dtqmaps[i], direction_e, mfem::get<i>(fops),
+                  integration_weights, scratch_mem);
+            }
+         }
+         else
+         {
+            map_field_to_quadrature_data(
+               directions_qp[i], dtqmaps[i], direction_e, mfem::get<i>(fops),
+               integration_weights);
+         }
+      }
+   });
+}
+
+}

examples/dfem/dfem_parametricspace.hpp

@@ -0,0 +1,104 @@
+#pragma once
+
+#include <mfem.hpp>
+
+namespace mfem
+{
+
+class ParametricSpace
+{
+
+public:
+   /// spatial_dim is the dimension of the spatial domain (e.g. 2 for 2D)
+   /// local_size is the size of the data on a single quadrature point
+   /// element_size is the size of the data on an element divided by vdim
+   /// total_size is the size of the data for all elements
+   ParametricSpace(int spatial_dim, int local_size, int element_size,
+                   int total_size) :
+      spatial_dim(spatial_dim),
+      local_size(local_size),
+      element_size(element_size),
+      total_size(total_size),
+      identity(total_size)
+   {
+      // dtq.ndof = (int)floor(pow(element_size, 1.0/spatial_dim) + 0.5);
+      dtq.ndof = element_size;
+      dtq.nqpt = dtq.ndof;
+   }
+
+   ParametricSpace(int local_size) :
+      local_size(local_size),
+      element_size(local_size),
+      total_size(local_size),
+      identity(local_size)
+   {
+      dtq.ndof = (int)floor(pow(element_size, 1.0/spatial_dim) + 0.5);
+      dtq.nqpt = dtq.ndof;
+   }
+
+   int Dimension() const
+   {
+      return spatial_dim;
+   }
+
+   int GetLocalSize() const
+   {
+      return local_size;
+   }
+
+   int GetElementSize() const
+   {
+      return element_size;
+   }
+
+   int GetTotalSize() const
+   {
+      return total_size;
+   }
+
+   const DofToQuad &GetDofToQuad() const
+   {
+      return dtq;
+   }
+
+   const Operator *GetProlongation() const
+   {
+      return &identity;
+   }
+
+   const Operator *GetRestriction() const
+   {
+      return &identity;
+   }
+
+private:
+   int spatial_dim;
+
+   // Hint for the local dimension. E.g. the size on the quadrature point or vdim.
+   int local_size;
+
+   // Size of the data on an element
+   int element_size;
+
+   int total_size;
+
+   IdentityOperator identity;
+
+   DofToQuad dtq;
+};
+
+class ParametricFunction : public Vector
+{
+public:
+   ParametricFunction(ParametricSpace &space) :
+      Vector(space.GetTotalSize()),
+      space(space)
+   {}
+
+   ParametricSpace &space;
+
+   using Vector::operator=;
+
+};
+
+}

examples/dfem/dfem_qfunction.hpp

@@ -0,0 +1,262 @@
+#pragma once
+#include "dfem_util.hpp"
+
+#ifdef MFEM_USE_ENZYME
+#include <enzyme/utils>
+#include <enzyme/enzyme>
+#endif
+
+namespace mfem
+{
+
+template <typename func_t, typename... arg_ts>
+inline auto qfunction_wrapper(const func_t &f, arg_ts &&...args)
+{
+   return f(args...);
+}
+
+template <typename T0, typename T1>
+MFEM_HOST_DEVICE
+void process_kf_arg(const T0 &, T1 &)
+{
+   static_assert(always_false<T0, T1>,
+                 "process_kf_arg not implemented for arg type");
+}
+
+template <typename T>
+MFEM_HOST_DEVICE
+void process_kf_arg(
+   const DeviceTensor<1, T> &u,
+   T &arg)
+{
+   arg = u(0);
+}
+
+template <typename T>
+MFEM_HOST_DEVICE
+void process_kf_arg(
+   const DeviceTensor<1, T> &u,
+   internal::tensor<T> &arg)
+{
+   arg(0) = u(0);
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::tensor<T, n> &arg)
+{
+   for (int i = 0; i < n; i++)
+   {
+      arg(i) = u(i);
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::tensor<T, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i) = u((i * m) + j);
+      }
+   }
+}
+
+template <typename arg_type>
+MFEM_HOST_DEVICE
+void process_kf_arg(const DeviceTensor<2> &u, arg_type &arg, int qp)
+{
+   const auto u_qp = Reshape(&u(0, qp), u.GetShape()[0]);
+   process_kf_arg(u_qp, arg);
+}
+
+template <size_t num_fields, typename kf_args>
+MFEM_HOST_DEVICE
+void process_kf_args(
+   const std::array<DeviceTensor<2>, num_fields> &u,
+   kf_args &args,
+   const int &qp)
+{
+   for_constexpr<mfem::tuple_size<kf_args>::value>([&](auto i)
+   {
+      process_kf_arg(u[i], mfem::get<i>(args), qp);
+      // out << mfem::get<i>(args) << ", ";
+   });
+}
+
+template <typename T0, typename T1> inline
+Vector process_kf_result(T0, T1)
+{
+   static_assert(always_false<T0, T1>,
+                 "process_kf_result not implemented for result type");
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const double &x)
+{
+   r(0) = x;
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<T> &x)
+{
+   r(0) = x(0);
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<T, n> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      r(i) = x(i);
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<T, n, m> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      for (size_t j = 0; j < m; j++)
+      {
+         r(i + n * j) = x(i, j);
+      }
+   }
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   double &arg)
+{
+   arg = u(0);
+}
+
+template <int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::tensor<double, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i) = u((i * m) + j);
+      }
+   }
+}
+
+template <typename qfunc_t, typename args_ts, size_t num_args>
+MFEM_HOST_DEVICE inline
+void apply_kernel(
+   DeviceTensor<1, double> &f_qp,
+   const qfunc_t &qfunc,
+   args_ts &args,
+   const std::array<DeviceTensor<2>, num_args> &u,
+   int qp)
+{
+   process_kf_args(u, args, qp);
+   process_kf_result(f_qp, mfem::get<0>(mfem::apply(qfunc, args)));
+}
+
+#ifdef MFEM_USE_ENZYME
+// Version for active function arguments only
+//
+// This is an Enzyme regression and can be removed in later versions.
+template <typename qfunc_t, typename arg_ts, std::size_t... Is,
+          typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme_indexed(qfunc_t &qfunc, arg_ts &&args,
+                                  arg_ts &&shadow_args,
+                                  std::index_sequence<Is...>,
+                                  inactive_arg_ts &&inactive_args,
+                                  std::index_sequence<>)
+{
+   using qf_return_t = typename create_function_signature<
+                       decltype(&qfunc_t::operator())>::type::return_t;
+   return __enzyme_fwddiff<qf_return_t>(
+             qfunction_wrapper<qfunc_t, decltype(mfem::get<Is>(args))...>, enzyme_const,
+             (void *)&qfunc, enzyme_dup, &mfem::get<Is>(args)..., enzyme_interleave,
+             &mfem::get<Is>(shadow_args)...);
+}
+
+// Interleave function arguments for enzyme
+template <typename qfunc_t, typename arg_ts, std::size_t... Is,
+          typename inactive_arg_ts, std::size_t... Js>
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme_indexed(qfunc_t &qfunc, arg_ts &&args,
+                                  arg_ts &&shadow_args,
+                                  std::index_sequence<Is...>,
+                                  inactive_arg_ts &&inactive_args,
+                                  std::index_sequence<Js...>)
+{
+   using qf_return_t = typename create_function_signature<
+                       decltype(&qfunc_t::operator())>::type::return_t;
+   return __enzyme_fwddiff<qf_return_t>(
+             qfunction_wrapper<qfunc_t, decltype(mfem::get<Is>(args))...,
+             decltype(mfem::get<Js>(inactive_args))...>,
+             enzyme_const, (void *)&qfunc, enzyme_dup, &mfem::get<Is>(args)...,
+             enzyme_const, &mfem::get<Js>(inactive_args)..., enzyme_interleave,
+             &mfem::get<Is>(shadow_args)...);
+}
+
+template <typename qfunc_t, typename arg_ts, typename inactive_arg_ts>
+MFEM_HOST_DEVICE inline
+auto fwddiff_apply_enzyme(qfunc_t &qfunc, arg_ts &&args,
+                          arg_ts &&shadow_args,
+                          inactive_arg_ts &&inactive_args)
+{
+   auto arg_indices = std::make_index_sequence<
+                      mfem::tuple_size<std::remove_reference_t<arg_ts>>::value> {};
+
+   auto inactive_arg_indices = std::make_index_sequence<
+                               mfem::tuple_size<std::remove_reference_t<inactive_arg_ts>>::value> {};
+
+   return fwddiff_apply_enzyme_indexed(qfunc, args, shadow_args, arg_indices,
+                                       inactive_args, inactive_arg_indices);
+}
+
+template <typename qfunc_t, typename arg_ts, size_t num_args>
+MFEM_HOST_DEVICE inline
+void apply_kernel_fwddiff_enzyme(
+   DeviceTensor<1, double> &f_qp,
+   qfunc_t &qfunc,
+   arg_ts &args,
+   arg_ts &shadow_args,
+   const std::array<DeviceTensor<2>, num_args> &u,
+   const std::array<DeviceTensor<2>, num_args> &v,
+   int qp_idx)
+{
+   // out << "\nargs: ";
+   process_kf_args(u, args, qp_idx);
+   // out << "\nshadow args: ";
+   process_kf_args(v, shadow_args, qp_idx);
+   // out << "\n";
+   process_kf_result(f_qp,
+                     mfem::get<0>(fwddiff_apply_enzyme(qfunc, args, shadow_args, mfem::tuple<> {})));
+}
+#endif // MFEM_USE_ENZYME
+
+} // namespace mfem

examples/dfem/dfem_qfunction_dual.hpp

@@ -0,0 +1,187 @@
+#pragma once
+#include "dfem_util.hpp"
+#include "dfem_qfunction.hpp"
+
+namespace mfem
+{
+
+MFEM_HOST_DEVICE
+template <typename T0, typename T1, typename T2>
+void process_kf_arg(const T0 &, const T1 &, T2 &)
+{
+   static_assert(always_false<T0, T1, T2>,
+                 "process_kf_arg not implemented for arg type");
+}
+
+template <typename T>
+MFEM_HOST_DEVICE
+void process_kf_arg(
+   const DeviceTensor<1, T> &u,
+   const DeviceTensor<1, T> &v,
+   T &arg)
+{
+   arg = u(0);
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::tensor<internal::dual<T, T>, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i).value = u((i * m) + j);
+      }
+   }
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   internal::dual<T, T> &arg)
+{
+   arg.value = u(0);
+}
+
+template <typename T>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::dual<T, T> &arg)
+{
+   arg.value = u(0);
+   arg.gradient = v(0);
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::tensor<internal::dual<T, T>, n> &arg)
+{
+   for (int i = 0; i < n; i++)
+   {
+      arg(i).value = u(i);
+      arg(i).gradient = v(i);
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<1> &u,
+   const DeviceTensor<1> &v,
+   internal::tensor<internal::dual<T, T>, n, m> &arg)
+{
+   for (int i = 0; i < m; i++)
+   {
+      for (int j = 0; j < n; j++)
+      {
+         arg(j, i).value = u((i * m) + j);
+         arg(j, i).gradient = v((i * m) + j);
+      }
+   }
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      r(i) = x(i).value;
+   }
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_kf_result(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n, m> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      for (size_t j = 0; j < m; j++)
+      {
+         r(i + n * j) = x(i, j).value;
+      }
+   }
+}
+
+template <typename arg_type>
+MFEM_HOST_DEVICE inline
+void process_kf_arg(
+   const DeviceTensor<2> &u,
+   const DeviceTensor<2> &v,
+   arg_type &arg,
+   const int &qp)
+{
+   const auto u_qp = Reshape(&u(0, qp), u.GetShape()[0]);
+   const auto v_qp = Reshape(&v(0, qp), v.GetShape()[0]);
+   process_kf_arg(u_qp, v_qp, arg);
+}
+
+template <size_t num_args, typename kf_args, std::size_t... Is>
+MFEM_HOST_DEVICE inline
+void process_kf_args(
+   const std::array<DeviceTensor<2>, num_args> &u,
+   const std::array<DeviceTensor<2>, num_args> &v,
+   kf_args &args,
+   const int &qp,
+   std::index_sequence<Is...>)
+{
+   (process_kf_arg(u[Is], v[Is], mfem::get<Is>(args), qp), ...);
+}
+
+template <typename T, int n, int m>
+MFEM_HOST_DEVICE inline
+void process_derivative_from_native_dual(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n, m> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      for (size_t j = 0; j < m; j++)
+      {
+         r(i + n * j) = x(i, j).gradient;
+      }
+   }
+}
+
+template <typename T, int n>
+MFEM_HOST_DEVICE inline
+void process_derivative_from_native_dual(
+   DeviceTensor<1, T> &r,
+   const internal::tensor<internal::dual<T, T>, n> &x)
+{
+   for (size_t i = 0; i < n; i++)
+   {
+      r(i) = x(i).gradient;
+   }
+}
+
+template <typename kf_t, typename kernel_arg_ts, size_t num_args>
+MFEM_HOST_DEVICE inline
+void apply_kernel_native_dual(
+   DeviceTensor<1, double> &f_qp,
+   const kf_t &kf,
+   kernel_arg_ts &args,
+   const std::array<DeviceTensor<2>, num_args> &u,
+   const std::array<DeviceTensor<2>, num_args> &v,
+   const int &qp_idx)
+{
+   process_kf_args(u, v, args, qp_idx,
+                   std::make_index_sequence<mfem::tuple_size<kernel_arg_ts>::value> {});
+   auto r = mfem::get<0>(mfem::apply(kf, args));
+   process_derivative_from_native_dual(f_qp, r);
+}
+
+} // namespace mfem

examples/dfem/dfem_refactor_action.hpp

@@ -0,0 +1,232 @@
+#pragma once
+
+#include "dfem_refactor.hpp"
+
+namespace mfem
+{
+
+template <typename element_operator_t, size_t num_fields>
+void DifferentiableOperator::instantiate_action(
+   element_operator_t element_operator, action_t &action)
+{
+   using entity_t = typename element_operator_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(
+                             fields,
+                             element_operator.inputs,
+                             std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(
+                              fields,
+                              element_operator.outputs,
+                              std::make_index_sequence<element_operator.num_outputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(element_operator.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(mesh);
+   const int num_entities = GetNumEntities<entity_t>(mesh);
+   const int num_qp = integration_rule.GetNPoints();
+
+   this->width = GetTrueVSize(fields[test_space_field_idx]);
+   size_t residual_lsize = GetVSize(fields[test_space_field_idx]);
+
+   // if constexpr (std::is_same_v<decltype(output_fop), One>)
+   // {
+   //    this->width = 1;
+   // }
+   // else
+   {
+      this->width = residual_lsize;
+   }
+
+   residual_l.SetSize(residual_lsize);
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(
+                          field,
+                          integration_rule,
+                          doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/mesh.Dimension()) + 0.5);
+
+   residual_e.SetSize(R->Height());
+
+   const int residual_size_on_qp = GetSizeOnQP<entity_t>(
+                                      mfem::get<hardcoded_output_idx>(element_operator.outputs),
+                                      fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(element_operator.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(element_operator.outputs, dtq,
+                                                    koutput_to_field);
+
+   // auto input_fops = create_bare_fops(element_operator.inputs);
+   // auto output_fops = create_bare_fops(element_operator.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>
+                         (element_operator.outputs).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(element_operator.inputs).size_on_qp /
+      mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(
+                              element_operator.inputs,
+                              std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto shmem_info =
+      get_shmem_info<entity_t, num_fields, element_operator.num_inputs, element_operator.num_outputs>
+      (input_dtq_maps,
+       output_dtq_maps,
+       fields,
+       num_entities,
+       element_operator.inputs,
+       num_qp,
+       input_size_on_qp,
+       residual_size_on_qp);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   print_shared_memory_info(shmem_info);
+
+   action = [=](const Vector &x, Vector &y) mutable
+   {
+      prolongation(solutions, x, solutions_l);
+
+      restriction<entity_t>(solutions, solutions_l, this->fields_e,
+                            element_dof_ordering);
+      restriction<entity_t>(parameters, parameters_l, this->fields_e,
+                            element_dof_ordering,
+                            solutions.size());
+
+      residual_e = 0.0;
+      auto ye = Reshape(residual_e.ReadWrite(), test_vdim, num_test_dof,
+                        num_entities);
+
+      auto wrapped_fields_e = wrap_fields(this->fields_e,
+                                          shmem_info.field_sizes,
+                                          num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, void *shmem)
+      {
+         // printf("\ne: %d\n", e);
+         // tic();
+         auto input_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::INPUT_DTQ],
+            shmem_info.input_dtq_sizes,
+            input_dtq_maps);
+
+         auto output_dtq_shmem = load_dtq_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::OUTPUT_DTQ],
+            shmem_info.output_dtq_sizes,
+            output_dtq_maps);
+
+         auto fields_shmem = load_field_mem(
+            shmem,
+            shmem_info.offsets[SharedMemory::Index::FIELD],
+            shmem_info.field_sizes,
+            kinput_to_field,
+            element_operator.inputs,
+            wrapped_fields_e,
+            e,
+         std::make_index_sequence<element_operator.num_inputs> {});
+
+         // These functions don't copy, they simply create a `DeviceTensor` object
+         // that points to correct chunks of the shared memory pool.
+         auto input_shmem = load_input_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::INPUT],
+                               shmem_info.input_sizes,
+                               num_qp);
+
+         auto residual_shmem = load_residual_mem(
+                                  shmem,
+                                  shmem_info.offsets[SharedMemory::Index::OUTPUT],
+                                  shmem_info.residual_size,
+                                  num_qp);
+
+         auto scratch_mem = load_scratch_mem(
+                               shmem,
+                               shmem_info.offsets[SharedMemory::Index::TEMP],
+                               shmem_info.temp_sizes);
+
+         MFEM_SYNC_THREAD;
+         // // printf("shmem load elapsed: %.1fus\n", toc() * 1e6);
+
+         // // tic();
+         map_fields_to_quadrature_data<TensorProduct>(
+            input_shmem, fields_shmem, input_dtq_shmem, element_operator.inputs, ir_weights,
+            scratch_mem,
+            std::make_index_sequence<element_operator.num_inputs> {});
+         // printf("interpolate elapsed: %.1fus\n", toc() * 1e6);
+
+         // // tic();
+         MFEM_FOREACH_THREAD(qx, x, q1d)
+         {
+            MFEM_FOREACH_THREAD(qy, y, q1d)
+            {
+               MFEM_FOREACH_THREAD(qz, z, q1d)
+               {
+                  const int q = qx + q1d * (qy + q1d * qz);
+                  auto qf_args = decay_tuple<typename element_operator_t::qf_param_ts> {};
+                  auto r = Reshape(&residual_shmem(0, q), residual_size_on_qp);
+                  apply_kernel(r, element_operator.qfunc, qf_args, input_shmem, q);
+               }
+            }
+         }
+         MFEM_SYNC_THREAD;
+         // // printf("qf elapsed: %.1fus\n", toc() * 1e6);
+
+         // // tic();
+         auto fhat = Reshape(&residual_shmem(0, 0), test_vdim, test_op_dim, num_qp);
+         auto y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         map_quadrature_data_to_fields<TensorProduct>(y, fhat,
+                                                      mfem::get<0>(element_operator.outputs),
+                                                      output_dtq_shmem[hardcoded_output_idx],
+                                                      scratch_mem);
+         // printf("integrate elapsed: %.1fus\n", toc() * 1e6);
+
+      }, num_entities, q1d, q1d, q1d, shmem_info.total_size, shmem_cache.ReadWrite());
+
+      if constexpr (std::is_same_v<decltype(output_fop), None<>>)
+      {
+         residual_l = y;
+      }
+      else
+      {
+         R->MultTranspose(residual_e, residual_l);
+      }
+
+      if constexpr (std::is_same_v<decltype(output_fop), None<>>)
+      {
+         y = residual_l;
+      }
+      // else if constexpr (std::is_same_v<decltype(output_fop), One>)
+      // {
+      //    double local_sum = residual_l.Sum();
+      //    MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM, mesh.GetComm());
+      //    MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      // }
+      else
+      {
+         get_prolongation(fields[test_space_field_idx])->MultTranspose(residual_l, y);
+      }
+   };
+}
+
+}

examples/dfem/dfem_refactor_derivatives.hpp

@@ -0,0 +1,132 @@
+#pragma once
+
+#include "dfem_refactor.hpp"
+
+template<typename T, T... Ints>
+void print_sequence(std::integer_sequence<T, Ints...>)
+{
+   ((std::cout << Ints << " "), ...);
+   std::cout << std::endl;
+}
+
+namespace mfem
+{
+
+template <
+   typename element_operator_t,
+   size_t num_solutions,
+   size_t num_parameters,
+   size_t derivative_idx>
+DerivativeOperator::DerivativeOperator(
+   element_operator_t element_operator,
+   const std::array<FieldDescriptor, num_solutions> &solutions,
+   const std::array<FieldDescriptor, num_parameters> &parameters,
+   const std::vector<FieldDescriptor> &fields,
+   ParMesh &mesh,
+   const IntegrationRule &integration_rule,
+   const ElementDofOrdering &element_dof_ordering,
+   const DofToQuad::Mode &doftoquad_mode,
+   std::integral_constant<size_t, derivative_idx>)
+{
+   direction = fields[derivative_idx];
+
+   size_t derivative_action_l_size = 0;
+   for (auto &s : solutions)
+   {
+      derivative_action_l_size += GetVSize(s);
+      this->width += GetTrueVSize(s);
+   }
+   this->height = derivative_action_l_size;
+   derivative_action_l.SetSize(derivative_action_l_size);
+
+   constexpr size_t num_fields = num_solutions + num_parameters;
+   using entity_t = typename element_operator_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(
+                             fields,
+                             element_operator.inputs,
+                             std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(
+                              fields,
+                              element_operator.outputs,
+                              std::make_index_sequence<element_operator.num_outputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(element_operator.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(mesh);
+   const int num_entities = GetNumEntities<entity_t>(mesh);
+   const int num_qp = integration_rule.GetNPoints();
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(
+                          field,
+                          integration_rule,
+                          doftoquad_mode));
+   }
+   const int q1d = (int)floor(pow(num_qp, 1.0/mesh.Dimension()) + 0.5);
+
+   derivative_action_e.SetSize(R->Height());
+
+   const int da_size_on_qp = GetSizeOnQP<entity_t>(
+                                mfem::get<hardcoded_output_idx>(element_operator.outputs),
+                                fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(element_operator.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(element_operator.outputs, dtq,
+                                                    koutput_to_field);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>
+                         (element_operator.outputs).vdim;
+   const int test_op_dim =
+      mfem::get<hardcoded_output_idx>(element_operator.inputs).size_on_qp /
+      mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim;
+   const int num_test_dof = R->Height() /
+                            mfem::get<hardcoded_output_idx>(element_operator.outputs).vdim /
+                            num_entities;
+
+   auto ir_weights = Reshape(integration_rule.GetWeights().Read(), num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(
+                              element_operator.inputs,
+                              std::make_index_sequence<element_operator.num_inputs> {});
+
+   auto input_is_dependent = std::get<derivative_idx>
+                             (element_operator.dependency_map);
+
+   constexpr bool with_derivatives = true;
+   auto shmem_info =
+      get_shmem_info<entity_t, num_fields, element_operator.num_inputs, element_operator.num_outputs>
+      (input_dtq_maps,
+       output_dtq_maps,
+       fields,
+       num_entities,
+       element_operator.inputs,
+       num_qp,
+       input_size_on_qp,
+       da_size_on_qp,
+       derivative_idx);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   print_shared_memory_info(shmem_info);
+
+   action_callback = [=](const Vector &x, Vector &y) mutable
+   {
+      restriction<entity_t>(direction, direction_l, direction_e,
+                            element_dof_ordering);
+
+   };
+}
+
+} // namespace mfem

examples/dfem/dfem_smm.hpp

@@ -0,0 +1,116 @@
+#pragma once
+
+#include <mfem.hpp>
+
+class SharedMemoryManager
+{
+private:
+   struct MemoryBlock
+   {
+      char* ptr;
+      int size;
+      bool used;
+   };
+
+   MFEM_HOST_DEVICE static const int MAX_BLOCKS = 16;
+   MFEM_HOST_DEVICE static MemoryBlock blocks[MAX_BLOCKS];
+   MFEM_HOST_DEVICE static int num_blocks;
+   MFEM_HOST_DEVICE static char* base_ptr;
+
+public:
+   MFEM_HOST_DEVICE static void init(void* shmem, int total_size)
+   {
+      base_ptr = static_cast<char*>(shmem);
+      num_blocks = 1;
+      blocks[0] = {base_ptr, total_size, false};
+   }
+
+   template<typename T>
+   MFEM_HOST_DEVICE static T* reserve(int n)
+   {
+      int size_bytes = n * sizeof(T);
+      for (int i = 0; i < num_blocks; ++i)
+      {
+         if (!blocks[i].used && blocks[i].size >= size_bytes)
+         {
+            blocks[i].used = true;
+            if (blocks[i].size > size_bytes)
+            {
+               // Split block
+               if (num_blocks < MAX_BLOCKS)
+               {
+                  blocks[num_blocks] = {blocks[i].ptr + size_bytes, blocks[i].size - size_bytes, false};
+                  ++num_blocks;
+                  blocks[i].size = size_bytes;
+               }
+            }
+            return reinterpret_cast<T*>(blocks[i].ptr);
+         }
+      }
+      return nullptr; // Allocation failed
+   }
+
+   MFEM_HOST_DEVICE static void release(void* ptr)
+   {
+      for (int i = 0; i < num_blocks; ++i)
+      {
+         if (blocks[i].ptr == ptr)
+         {
+            blocks[i].used = false;
+            return;
+         }
+      }
+   }
+
+   MFEM_HOST_DEVICE static void release_and_try_merge(void* ptr)
+   {
+      for (int i = 0; i < num_blocks; ++i)
+      {
+         if (blocks[i].ptr == ptr)
+         {
+            blocks[i].used = false;
+            merge_adjacent_free_blocks();
+            return;
+         }
+      }
+   }
+
+private:
+   MFEM_HOST_DEVICE static void merge_adjacent_free_blocks()
+   {
+      // Simple bubble sort for simplicity (can be optimized)
+      for (int i = 0; i < num_blocks - 1; ++i)
+      {
+         for (int j = 0; j < num_blocks - i - 1; ++j)
+         {
+            if (blocks[j].ptr > blocks[j + 1].ptr)
+            {
+               MemoryBlock temp = blocks[j];
+               blocks[j] = blocks[j + 1];
+               blocks[j + 1] = temp;
+            }
+         }
+      }
+
+      for (int i = 0; i < num_blocks - 1; ++i)
+      {
+         if (!blocks[i].used && !blocks[i + 1].used)
+         {
+            blocks[i].size += blocks[i + 1].size;
+            for (int j = i + 1; j < num_blocks - 1; ++j)
+            {
+               blocks[j] = blocks[j + 1];
+            }
+            --num_blocks;
+            --i;
+         }
+      }
+   }
+};
+
+MFEM_HOST_DEVICE SharedMemoryManager::MemoryBlock
+SharedMemoryManager::blocks[SharedMemoryManager::MAX_BLOCKS];
+
+MFEM_HOST_DEVICE int SharedMemoryManager::num_blocks;
+
+MFEM_HOST_DEVICE char* SharedMemoryManager::base_ptr;

examples/dfem/dfem_test_macro.hpp

@@ -0,0 +1,39 @@
+#pragma once
+#include "dfem_refactor.hpp"
+
+#define DFEM_TEST_MAIN(function)                                               \
+   int main(int argc, char* argv[])                                            \
+   {                                                                           \
+      Mpi::Init();                                                             \
+                                                                               \
+      const char* device_config = "cpu";                                       \
+      const char* mesh_file = "../data/ref-square.mesh";                       \
+      int polynomial_order = 1;                                                \
+      int ir_order = 2;                                                        \
+      int refinements = 0;                                                     \
+                                                                               \
+      OptionsParser args(argc, argv);                                          \
+      args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");         \
+      args.AddOption(&polynomial_order, "-o", "--order", "");                  \
+      args.AddOption(&refinements, "-r", "--r", "");                           \
+      args.AddOption(&ir_order, "-iro", "--iro", "");                          \
+      args.AddOption(&device_config, "-d", "--device",                         \
+         "Device configuration string, see Device::Configure().");             \
+      args.ParseCheck();                                                       \
+                                                                               \
+      Device device(device_config);                                            \
+      if (Mpi::Root() == 0)                                                    \
+      {                                                                        \
+         device.Print();                                                       \
+      }                                                                        \
+                                                                               \
+      out << std::setprecision(12);                                            \
+                                                                               \
+      int ret;                                                                 \
+                                                                               \
+      ret = function(mesh_file, refinements, polynomial_order);                \
+      out << #function;                                                        \
+      ret ? out << " FAILURE\n" : out << " OK\n";                              \
+                                                                               \
+      return ret;                                                              \
+   }\

examples/dfem/noisy.hpp

@@ -0,0 +1,130 @@
+// SPDX-ArtifactOfProjectName:     noisy
+// SPDX-ArtifactOfProjectHomePage: https://github.com/VincentZalzal/noisy
+// SPDX-FileCopyrightText:         Copyright 2024 Vincent Zalzal
+// SPDX-License-Identifier:        MIT
+
+#pragma once
+
+#include <iomanip>
+#include <iostream>
+
+namespace vz {
+
+struct Counters {
+    unsigned m_def_ctor    = 0;
+    unsigned m_copy_ctor   = 0;
+    unsigned m_move_ctor   = 0;
+    unsigned m_copy_assign = 0;
+    unsigned m_move_assign = 0;
+    unsigned m_dtor        = 0;
+
+    void reset() {
+        *this = {};
+    }
+
+    bool leaks() const {
+        return m_def_ctor + m_copy_ctor + m_move_ctor != m_dtor;
+    }
+
+    friend std::ostream& operator<<(std::ostream& os, const Counters& c) {
+        stream_counter(os, "Default constructor count: ", c.m_def_ctor   );
+        stream_counter(os, "Copy constructor count:    ", c.m_copy_ctor  );
+        stream_counter(os, "Move constructor count:    ", c.m_move_ctor  );
+        stream_counter(os, "Copy assignment count:     ", c.m_copy_assign);
+        stream_counter(os, "Move assignment count:     ", c.m_move_assign);
+        stream_counter(os, "Destructor count:          ", c.m_dtor       );
+        return os;
+    }
+
+    friend bool operator==(const Counters& lhs, const Counters& rhs) {
+        return
+            lhs.m_def_ctor    == rhs.m_def_ctor    &&
+            lhs.m_copy_ctor   == rhs.m_copy_ctor   &&
+            lhs.m_move_ctor   == rhs.m_move_ctor   &&
+            lhs.m_copy_assign == rhs.m_copy_assign &&
+            lhs.m_move_assign == rhs.m_move_assign &&
+            lhs.m_dtor        == rhs.m_dtor        ;
+    }
+
+    friend bool operator!=(const Counters& lhs, const Counters& rhs) { return !(lhs == rhs); }
+
+private:
+    static void stream_counter(std::ostream& os, const char* msg, unsigned value) {
+        if (value != 0)
+            os << msg << std::setw(2) << value << '\n';
+    }
+};
+
+namespace detail {
+
+struct Globals {
+    ~Globals() {
+        if (m_verbose)
+            std::cout << "\n===== Noisy counters =====\n" << m_counters;
+    }
+
+    Counters m_counters;
+    unsigned m_next_id = 0;
+    bool     m_verbose = true;
+};
+
+}
+
+class Noisy {
+private:
+    static detail::Globals& globals() {
+        static detail::Globals s_globals;
+        return s_globals;
+    }
+
+public:
+    static Counters& counters() { return globals().m_counters; }
+    static void set_verbose(bool verbose) { globals().m_verbose = verbose; }
+
+    Noisy() {
+        if (globals().m_verbose)
+            std::cout << *this << ": default constructor\n";
+        globals().m_counters.m_def_ctor++;
+    }
+
+    Noisy(const Noisy& other) {
+        if (globals().m_verbose)
+            std::cout << *this << ": copy constructor from " << other << '\n';
+        globals().m_counters.m_copy_ctor++;
+    }
+
+    Noisy(Noisy&& other) noexcept {
+        if (globals().m_verbose)
+            std::cout << *this << ": move constructor from " << other << '\n';
+        globals().m_counters.m_move_ctor++;
+    }
+
+    ~Noisy() {
+        if (globals().m_verbose)
+            std::cout << *this << ": destructor\n";
+        globals().m_counters.m_dtor++;
+    }
+
+    Noisy& operator=(const Noisy& other) {
+        if (globals().m_verbose)
+            std::cout << *this << ": copy assignment from " << other << '\n';
+        globals().m_counters.m_copy_assign++;
+        return *this;
+    }
+
+    Noisy& operator=(Noisy&& other) noexcept {
+        if (globals().m_verbose)
+            std::cout << *this << ": move assignment from " << other << '\n';
+        globals().m_counters.m_move_assign++;
+        return *this;
+    }
+
+    unsigned id() const { return m_id; }
+
+    friend std::ostream& operator<<(std::ostream& os, const Noisy& noisy) { return os << "Noisy(" << std::setw(2) << noisy.m_id << ')'; }
+
+private:
+    unsigned m_id = globals().m_next_id++;
+};
+
+}

examples/dfem/old_derivative_action_callback.icc

@@ -0,0 +1,188 @@
+{
+   using entity_t = typename kernel_t::entity_t;
+
+   auto kinput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+   kernel.inputs, std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto koutput_to_field = create_descriptors_to_fields_map<entity_t>(op.fields,
+   kernel.outputs, std::make_index_sequence<kernel.num_koutputs> {});
+
+   constexpr int hardcoded_output_idx = 0;
+   const int test_space_field_idx = koutput_to_field[hardcoded_output_idx];
+
+   const Operator *R = get_restriction<entity_t>(op.fields[test_space_field_idx],
+                                                 element_dof_ordering);
+
+   auto output_fop = mfem::get<hardcoded_output_idx>(kernel.outputs);
+
+   const int num_elements = GetNumEntities<Entity::Element>(op.mesh);
+   const int num_entities = GetNumEntities<entity_t>(op.mesh);
+   const int num_qp = op.integration_rule.GetNPoints();
+
+   // assume only a single element type for now
+   std::vector<const DofToQuad*> dtq;
+   for (const auto &field : op.fields)
+   {
+      dtq.emplace_back(GetDofToQuad<entity_t>(field, op.integration_rule,
+                                              doftoquad_mode));
+   }
+   const int q1d = dtq[0]->nqpt;
+
+   derivative_action_e.SetSize(R->Height());
+
+   const int da_size_on_qp = GetSizeOnQP<entity_t>(
+      mfem::get<hardcoded_output_idx>(kernel.outputs),
+      op.fields[test_space_field_idx]);
+
+   auto input_dtq_maps = create_dtq_maps<entity_t>(kernel.inputs, dtq,
+                                                   kinput_to_field);
+   auto output_dtq_maps = create_dtq_maps<entity_t>(kernel.outputs, dtq,
+                                                    koutput_to_field);
+
+   auto input_fops = create_bare_fops(kernel.inputs);
+   auto output_fops = create_bare_fops(kernel.outputs);
+
+   const int test_vdim = mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int test_op_dim =
+   mfem::get<hardcoded_output_idx>(output_fops).size_on_qp /
+   mfem::get<hardcoded_output_idx>(output_fops).vdim;
+   const int num_test_dof = R->Height() /
+   mfem::get<hardcoded_output_idx>(output_fops).vdim /
+   num_entities;
+
+   auto ir_weights = Reshape(this->op.integration_rule.GetWeights().Read(),
+                             num_qp);
+
+   auto input_size_on_qp = get_input_size_on_qp(kernel.inputs,
+   std::make_index_sequence<kernel.num_kinputs> {});
+
+   auto shmem_info = get_shmem_info<entity_t>(input_dtq_maps,
+                                              output_dtq_maps,
+                                              op.fields,
+                                              num_entities,
+                                              kernel.inputs,
+                                              num_qp,
+                                              input_size_on_qp,
+                                              da_size_on_qp);
+
+   Vector shmem_cache(shmem_info.total_size);
+
+   func = [=](Vector &ye_mem) mutable
+   {
+      restriction<entity_t>(direction, direction_l, direction_e,
+                            op.element_dof_ordering, derivative_idx);
+
+      // Check which qf inputs are dependent on the dependent variable
+      std::array<bool, kernel.num_kinputs> kinput_is_dependent;
+      bool no_qfinput_is_dependent = true;
+      for (int i = 0; i < kinput_is_dependent.size(); i++)
+      {
+         if (kinput_to_field[i] == derivative_idx)
+         {
+            no_qfinput_is_dependent = false;
+            kinput_is_dependent[i] = true;
+            // out << "function input " << i << " is dependent on "
+            //     << op.fields[kinput_to_field[i]].field_label << "\n";
+         }
+         else
+         {
+            kinput_is_dependent[i] = false;
+         }
+      }
+
+      if (no_qfinput_is_dependent)
+      {
+         return;
+      }
+
+      // auto kernel_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+      // auto kernel_shadow_args = decay_tuple<typename kernel_t::kf_param_ts> {};
+
+      // DeviceTensor<1, const double> integration_weights(
+      //    this->op.integration_rule.GetWeights().Read(), num_qp);
+
+      // Vector zero;
+      // GeometricFactorMaps geometric_factors
+      // {
+      //    DeviceTensor<3, const double>(zero.Read(), 0, 0, 0)
+      // };
+
+      // // Fields interpolated to the quadrature points in the order of
+      // // kernel function arguments
+      // auto input_qp = map_inputs_to_memory(input_qp_mem, num_qp,
+      //                                      kernel.inputs,
+      //                                      std::make_index_sequence<kernel.num_kinputs> {});
+
+      // auto directions_qp = map_inputs_to_memory(directions_qp_mem, num_qp,
+      //                                           kernel.inputs,
+      //                                           std::make_index_sequence<kernel.num_kinputs> {});
+
+      // constexpr int fixed_output_idx = 0;
+      // auto Bv = output_dtq_maps[fixed_output_idx];
+      // auto [num_test_qp, test_op_dim, num_test_dof] = Bv.GetShape();
+      // const int test_vdim = mfem::get<0>(kernel.outputs).vdim;
+      // DeviceTensor<3> ye = Reshape(ye_mem.ReadWrite(), num_test_dof, test_vdim, num_entities);
+
+      forall([=] MFEM_HOST_DEVICE (int e, double *shmem)
+      {
+         // map_fields_to_quadrature_data(
+         //    input_qp, e, this->fields_e,
+         //    kinput_to_field, input_dtq_maps,
+         //    integration_weights, geometric_factors, kernel.inputs,
+         //    std::make_index_sequence<kernel.num_kinputs> {});
+
+         // map_fields_to_quadrature_data_conditional(
+         //    directions_qp, e,
+         //    directions_e, kinput_to_field,
+         //    input_dtq_maps,
+         //    integration_weights,
+         //    geometric_factors,
+         //    kinput_is_dependent,
+         //    kernel.inputs,
+         //    std::make_index_sequence<kernel.num_kinputs> {});
+
+         // for (int qp = 0; qp < num_qp; qp++)
+         // {
+         //    auto f_qp = apply_kernel_fwddiff_enzyme(
+         //                   kernel.func,
+         //                   kernel_args,
+         //                   input_qp,
+         //                   kernel_shadow_args,
+         //                   directions_qp,
+         //                   qp);
+
+         //    auto r_qp = Reshape(&da_qp(0, qp, e), da_size_on_qp);
+         //    for (int i = 0; i < da_size_on_qp; i++)
+         //    {
+         //       r_qp(i) = f_qp(i);
+         //    }
+         // }
+
+         // DeviceTensor<3> fhat = Reshape(&da_qp(0, 0, e), test_vdim, test_op_dim, num_qp);
+         // DeviceTensor<2> y = Reshape(&ye(0, 0, e), num_test_dof, test_vdim);
+         // map_quadrature_data_to_fields(y, fhat,
+         //                               output_fop,
+         //                               output_dtq_maps[hardcoded_output_idx]);
+      }, num_entities, q1d, q1d, 1, shmem_info.total_size, shmem_cache.GetData());
+
+      R->MultTranspose(ye_mem, derivative_action_l);
+   };
+
+   if constexpr (std::is_same_v<decltype(output_fop), One>)
+   {
+      prolongation_transpose = [&](Vector &r_local, Vector &y)
+      {
+         double local_sum = r_local.Sum();
+         MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM,
+                       op.mesh.GetComm());
+         MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
+      };
+   }
+   else
+   {
+      auto P = get_prolongation(op.fields[test_space_field_idx]);
+      prolongation_transpose = [P](Vector &r_l, Vector &y)
+      {
+         P->MultTranspose(r_l, y);
+      };
+   }

examples/dfem/shared_mem_calc.md

@@ -0,0 +1,49 @@
+* Calculate shared memory requirements
+* Interpolation and integration
+---
+* If grad involved, need B and G
+* Fit largest field, depends on polynomial order (#dofs)
+  -> vdim is irrelevant
+* Temporaries for each sum
+  - DDQ (d1d x d1d x q1d) x 2 -> DDQ0, DDQ1
+  - DQQ (d1d x q1d x q1d) x 3 -> DQQ0, DQQ1, DQQ2
+  - QQQ (q1d x q1d x q1d) x 3 -> QQQ0, QQQ1, QQQ2
+
+We need the following combinations at the same time
+(1) DDQ0 + DDQ1 + DQQ0 + DQQ1 + DQQ2
+(2) DQQ0 + DQQ1 + DQQ2 + QQQ0 + QQQ1 + QQQ2
+(3) QQQ0 + QQQ1 + QQQ2 + QQD0 + QQD1 + QQD2
+(4) QQD0 + QQD1 + QQD2 + QDD0 + QDD1 + QDD2
+
+Allocate largest memory footprint from 2, 3 or 4 and
+add memory footprint of fields and B/G.
+
+Annotations with NR and R mean "not reusable" and
+"reusable", respectively. This means the memory location is
+reused for _all_ e.g. interpolation of a value etc.
+
+----
+For the action of nonlinear diffusion in 2D we have
+(rho * |u|^2 \nabla u, \nabla v)
+
+* Load
+RHO (D x D) | R (after interpolation)
+U (D x D x VDIM) | R (after interpolation)
+B (Q x D) | NR
+G (Q x D) | NR
+
+* Interpolate Value
+Temporary (Q x D) | R
+R (Q x Q) | NR
+U (Q x Q x VDIM) | NR
+
+* Interpolate Grad
+Temporaries (Q x D) + (Q x D) | R
+U (Q x Q x DIM x VDIM) | NR
+
+Quadrature point function
+-> purely thread local
+
+* Integrate Grad
+R | temp from Interpolation
+R | U from Load

examples/dfem/tuple.hpp

@@ -0,0 +1,845 @@
+// This is serac's tuple implementation
+
+#pragma once
+
+#include "general/backends.hpp"
+#include <utility>
+#include <mfem.hpp>
+#include <tuple>
+
+namespace mfem
+{
+
+/**
+ * @tparam T the types stored in the tuple
+ * @brief This is a class that mimics most of std::tuple's interface,
+ * except that it is usable in CUDA kernels and admits some arithmetic operator overloads.
+ *
+ * see https://en.cppreference.com/w/cpp/utility/tuple for more information about std::tuple
+ */
+template <typename... T>
+struct tuple
+{
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ */
+template <typename T0>
+struct tuple<T0>
+{
+   T0 v0;  ///< The first member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ */
+template <typename T0, typename T1>
+struct tuple<T0, T1>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2>
+struct tuple<T0, T1, T2>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3>
+struct tuple<T0, T1, T2, T3>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4>
+struct tuple<T0, T1, T2, T3, T4>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ * @tparam T5 The sixth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5>
+struct tuple<T0, T1, T2, T3, T4, T5>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ * @tparam T5 The sixth type stored in the tuple
+ * @tparam T6 The seventh type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6>
+struct tuple<T0, T1, T2, T3, T4, T5, T6>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+   T6 v6;  ///< The seventh member of the tuple
+};
+
+/**
+ * @brief Type that mimics std::tuple
+ *
+ * @tparam T0 The first type stored in the tuple
+ * @tparam T1 The second type stored in the tuple
+ * @tparam T2 The third type stored in the tuple
+ * @tparam T3 The fourth type stored in the tuple
+ * @tparam T4 The fifth type stored in the tuple
+ * @tparam T5 The sixth type stored in the tuple
+ * @tparam T6 The seventh type stored in the tuple
+ * @tparam T7 The eighth type stored in the tuple
+ */
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7>
+struct tuple<T0, T1, T2, T3, T4, T5, T6, T7>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+   T6 v6;  ///< The seventh member of the tuple
+   T7 v7;  ///< The eighth member of the tuple
+};
+
+template <typename T0, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8>
+struct tuple<T0, T1, T2, T3, T4, T5, T6, T7, T8>
+{
+   T0 v0;  ///< The first member of the tuple
+   T1 v1;  ///< The second member of the tuple
+   T2 v2;  ///< The third member of the tuple
+   T3 v3;  ///< The fourth member of the tuple
+   T4 v4;  ///< The fifth member of the tuple
+   T5 v5;  ///< The sixth member of the tuple
+   T6 v6;  ///< The seventh member of the tuple
+   T7 v7;  ///< The eighth member of the tuple
+   T8 v8;
+};
+
+/**
+ * @brief Class template argument deduction rule for tuples
+ * @tparam T The variadic template parameter for tuple types
+ */
+template <typename... T>
+MFEM_HOST_DEVICE
+tuple(T...) -> tuple<T...>;
+
+/**
+ * @brief helper function for combining a list of values into a tuple
+ * @tparam T types of the values to be tuple-d
+ * @param args the actual values to be put into a tuple
+ */
+template <typename... T>
+MFEM_HOST_DEVICE tuple<T...> make_tuple(const T&... args)
+{
+   return tuple<T...> {args...};
+}
+
+template <class... Types>
+struct tuple_size
+{
+};
+
+template <class... Types>
+struct tuple_size<mfem::tuple<Types...>> :
+                                      std::integral_constant<std::size_t, sizeof...(Types)>
+{
+};
+
+/**
+ * @tparam i the tuple index to access
+ * @tparam T the types stored in the tuple
+ * @brief return a reference to the ith tuple entry
+ */
+template <int i, typename... T>
+MFEM_HOST_DEVICE constexpr auto& get(tuple<T...>& values)
+{
+   static_assert(i < sizeof...(T), "");
+   if constexpr (i == 0)
+   {
+      return values.v0;
+   }
+   if constexpr (i == 1)
+   {
+      return values.v1;
+   }
+   if constexpr (i == 2)
+   {
+      return values.v2;
+   }
+   if constexpr (i == 3)
+   {
+      return values.v3;
+   }
+   if constexpr (i == 4)
+   {
+      return values.v4;
+   }
+   if constexpr (i == 5)
+   {
+      return values.v5;
+   }
+   if constexpr (i == 6)
+   {
+      return values.v6;
+   }
+   if constexpr (i == 7)
+   {
+      return values.v7;
+   }
+   if constexpr (i == 8)
+   {
+      return values.v8;
+   }
+}
+
+/**
+ * @tparam i the tuple index to access
+ * @tparam T the types stored in the tuple
+ * @brief return a copy of the ith tuple entry
+ */
+template <int i, typename... T>
+MFEM_HOST_DEVICE constexpr const auto& get(const tuple<T...>& values)
+{
+   static_assert(i < sizeof...(T), "");
+   if constexpr (i == 0)
+   {
+      return values.v0;
+   }
+   if constexpr (i == 1)
+   {
+      return values.v1;
+   }
+   if constexpr (i == 2)
+   {
+      return values.v2;
+   }
+   if constexpr (i == 3)
+   {
+      return values.v3;
+   }
+   if constexpr (i == 4)
+   {
+      return values.v4;
+   }
+   if constexpr (i == 5)
+   {
+      return values.v5;
+   }
+   if constexpr (i == 6)
+   {
+      return values.v6;
+   }
+   if constexpr (i == 7)
+   {
+      return values.v7;
+   }
+   if constexpr (i == 8)
+   {
+      return values.v8;
+   }
+}
+
+/**
+ * @brief a function intended to be used for extracting the ith type from a tuple.
+ *
+ * @note type<i>(my_tuple) returns a value, whereas get<i>(my_tuple) returns a reference
+ *
+ * @tparam i the index of the tuple to query
+ * @tparam T the types stored in the tuple
+ * @param values the tuple of values
+ * @return a copy of the ith entry of the input
+ */
+template <int i, typename... T>
+MFEM_HOST_DEVICE constexpr auto type(const tuple<T...>& values)
+{
+   static_assert(i < sizeof...(T), "");
+   if constexpr (i == 0)
+   {
+      return values.v0;
+   }
+   if constexpr (i == 1)
+   {
+      return values.v1;
+   }
+   if constexpr (i == 2)
+   {
+      return values.v2;
+   }
+   if constexpr (i == 3)
+   {
+      return values.v3;
+   }
+   if constexpr (i == 4)
+   {
+      return values.v4;
+   }
+   if constexpr (i == 5)
+   {
+      return values.v5;
+   }
+   if constexpr (i == 6)
+   {
+      return values.v6;
+   }
+   if constexpr (i == 7)
+   {
+      return values.v7;
+   }
+   if constexpr (i == 8)
+   {
+      return values.v8;
+   }
+}
+
+/**
+ * @brief A helper function for the + operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple sum
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto plus_helper(const tuple<S...>& x,
+                                            const tuple<T...>& y,
+                                            std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) + get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise sum of x and y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator+(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return plus_helper(x, y,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the += operator of tuples
+ *
+ * @tparam T the types stored in the tuples x and y
+ * @tparam i integer sequence used to index the tuples
+ * @param x tuple of values to be incremented
+ * @param y tuple of increment values
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr void plus_equals_helper(tuple<T...>& x,
+                                                   const tuple<T...>& y,
+                                                   std::integer_sequence<int, i...>)
+{
+   ((get<i>(x) += get<i>(y)), ...);
+}
+
+/**
+ * @tparam T the types stored in the tuples x and y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief add values contained in y, to the tuple x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator+=(tuple<T...>& x,
+                                           const tuple<T...>& y)
+{
+   return plus_equals_helper(x, y,
+                             std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the -= operator of tuples
+ *
+ * @tparam T the types stored in the tuples x and y
+ * @tparam i integer sequence used to index the tuples
+ * @param x tuple of values to be subracted from
+ * @param y tuple of values to subtract from x
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr void minus_equals_helper(tuple<T...>& x,
+                                                    const tuple<T...>& y,
+                                                    std::integer_sequence<int, i...>)
+{
+   ((get<i>(x) -= get<i>(y)), ...);
+}
+
+/**
+ * @tparam T the types stored in the tuples x and y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief add values contained in y, to the tuple x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator-=(tuple<T...>& x,
+                                           const tuple<T...>& y)
+{
+   return minus_equals_helper(x, y,
+                              std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the - operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple difference
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto minus_helper(const tuple<S...>& x,
+                                             const tuple<T...>& y,
+                                             std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) - get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise difference of x and y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator-(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return minus_helper(x, y,
+                       std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the - operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @return the returned tuple difference
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto unary_minus_helper(const tuple<T...>& x,
+                                                   std::integer_sequence<int, i...>)
+{
+   return tuple{-get<i>(x)...};
+}
+
+/**
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @brief return a tuple of values defined by applying the unary minus operator to each element of x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator-(const tuple<T...>& x)
+{
+   return unary_minus_helper(x,
+                             std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the / operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple ratio
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto div_helper(const tuple<S...>& x,
+                                           const tuple<T...>& y,
+                                           std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) / get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise division of x by y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator/(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return div_helper(x, y,
+                     std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the / operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a the constant numerator
+ * @return the returned tuple ratio
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto div_helper(const double a,
+                                           const tuple<T...>& x, std::integer_sequence<int, i...>)
+{
+   return tuple{a / get<i>(x)...};
+}
+
+/**
+ * @brief A helper function for the / operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a the constant denomenator
+ * @return the returned tuple ratio
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto div_helper(const tuple<T...>& x,
+                                           const double a, std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) / a...};
+}
+
+/**
+ * @tparam T the types stored in the tuple x
+ * @param a the numerator
+ * @param x a tuple of denominator values
+ * @brief return a tuple of values defined by division of a by the elements of x
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator/(const double a, const tuple<T...>& x)
+{
+   return div_helper(a, x,
+                     std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of numerator values
+ * @param a a denominator
+ * @brief return a tuple of values defined by elementwise division of x by a
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator/(const tuple<T...>& x, const double a)
+{
+   return div_helper(x, a,
+                     std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief A helper function for the * operator of tuples
+ *
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param y tuple of values
+ * @return the returned tuple product
+ */
+template <typename... S, typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto mult_helper(const tuple<S...>& x,
+                                            const tuple<T...>& y,
+                                            std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) * get<i>(y)...};
+}
+
+/**
+ * @tparam S the types stored in the tuple x
+ * @tparam T the types stored in the tuple y
+ * @param x a tuple of values
+ * @param y a tuple of values
+ * @brief return a tuple of values defined by elementwise multiplication of x and y
+ */
+template <typename... S, typename... T>
+MFEM_HOST_DEVICE constexpr auto operator*(const tuple<S...>& x,
+                                          const tuple<T...>& y)
+{
+   static_assert(sizeof...(S) == sizeof...(T));
+   return mult_helper(x, y,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(S))>());
+}
+
+/**
+ * @brief A helper function for the * operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a a constant multiplier
+ * @return the returned tuple product
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto mult_helper(const double a,
+                                            const tuple<T...>& x, std::integer_sequence<int, i...>)
+{
+   return tuple{a * get<i>(x)...};
+}
+
+/**
+ * @brief A helper function for the * operator of tuples
+ *
+ * @tparam T the types stored in the tuple y
+ * @tparam i The integer sequence to i
+ * @param x tuple of values
+ * @param a a constant multiplier
+ * @return the returned tuple product
+ */
+template <typename... T, int... i>
+MFEM_HOST_DEVICE constexpr auto mult_helper(const tuple<T...>& x,
+                                            const double a, std::integer_sequence<int, i...>)
+{
+   return tuple{get<i>(x) * a...};
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @param a a scaling factor
+ * @param x the tuple object
+ * @brief multiply each component of x by the value a on the left
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator*(const double a, const tuple<T...>& x)
+{
+   return mult_helper(a, x,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @param x the tuple object
+ * @param a a scaling factor
+ * @brief multiply each component of x by the value a on the right
+ */
+template <typename... T>
+MFEM_HOST_DEVICE constexpr auto operator*(const tuple<T...>& x, const double a)
+{
+   return mult_helper(x, a,
+                      std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @tparam i a list of indices used to acces each element of the tuple
+ * @param out the ostream to write the output to
+ * @param A the tuple of values
+ * @brief helper used to implement printing a tuple of values
+ */
+template <typename... T, std::size_t... i>
+auto& print_helper(std::ostream& out, const mfem::tuple<T...>& A,
+                   std::integer_sequence<size_t, i...>)
+{
+   out << "tuple{";
+   (..., (out << (i == 0 ? "" : ", ") << mfem::get<i>(A)));
+   out << "}";
+   return out;
+}
+
+/**
+ * @tparam T the types stored in the tuple
+ * @param out the ostream to write the output to
+ * @param A the tuple of values
+ * @brief print a tuple of values
+ */
+template <typename... T>
+auto& operator<<(std::ostream& out, const mfem::tuple<T...>& A)
+{
+   return print_helper(out, A, std::make_integer_sequence<size_t, sizeof...(T)>());
+}
+
+/**
+ * @brief A helper to apply a lambda to a tuple
+ *
+ * @tparam lambda The functor type
+ * @tparam T The tuple types
+ * @tparam i The integer sequence to i
+ * @param f The functor to apply to the tuple
+ * @param args The input tuple
+ * @return The functor output
+ */
+template <typename lambda, typename... T, int... i>
+MFEM_HOST_DEVICE auto apply_helper(lambda f, tuple<T...>& args,
+                                   std::integer_sequence<int, i...>)
+{
+   return f(get<i>(args)...);
+}
+
+/**
+ * @tparam lambda a callable type
+ * @tparam T the types of arguments to be passed in to f
+ * @param f the callable object
+ * @param args a tuple of arguments
+ * @brief a way of passing an n-tuple to a function that expects n separate arguments
+ *
+ *   e.g. foo(bar, baz) is equivalent to apply(foo, mfem::tuple(bar,baz));
+ */
+template <typename lambda, typename... T>
+MFEM_HOST_DEVICE auto apply(lambda f, tuple<T...>& args)
+{
+   return apply_helper(f, std::move(args),
+                       std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @overload
+ */
+template <typename lambda, typename... T, int... i>
+MFEM_HOST_DEVICE auto apply_helper(lambda f, const tuple<T...>& args,
+                                   std::integer_sequence<int, i...>)
+{
+   return f(get<i>(args)...);
+}
+
+/**
+ * @tparam lambda a callable type
+ * @tparam T the types of arguments to be passed in to f
+ * @param f the callable object
+ * @param args a tuple of arguments
+ * @brief a way of passing an n-tuple to a function that expects n separate arguments
+ *
+ *   e.g. foo(bar, baz) is equivalent to apply(foo, mfem::tuple(bar,baz));
+ */
+template <typename lambda, typename... T>
+MFEM_HOST_DEVICE auto apply(lambda f, const tuple<T...>& args)
+{
+   return apply_helper(f, std::move(args),
+                       std::make_integer_sequence<int, static_cast<int>(sizeof...(T))>());
+}
+
+/**
+ * @brief a struct used to determine the type at index I of a tuple
+ *
+ * @note see: https://en.cppreference.com/w/cpp/utility/tuple/tuple_element
+ *
+ * @tparam I the index of the desired type
+ * @tparam T a tuple of different types
+ */
+template <size_t I, class T>
+struct tuple_element;
+
+// recursive case
+/// @overload
+template <size_t I, class Head, class... Tail>
+struct tuple_element<I, tuple<Head, Tail...>> : tuple_element<I - 1,
+                                              tuple<Tail...>>
+{
+};
+
+// base case
+/// @overload
+template <class Head, class... Tail>
+struct tuple_element<0, tuple<Head, Tail...>>
+{
+   using type = Head;  ///< the type at the specified index
+};
+
+/**
+ * @brief Trait for checking if a type is a @p mfem::tuple
+ */
+template <typename T>
+struct is_tuple : std::false_type
+{
+};
+
+/// @overload
+template <typename... T>
+struct is_tuple<mfem::tuple<T...>> : std::true_type
+{
+};
+
+/**
+ * @brief Trait for checking if a type if a @p mfem::tuple containing only @p mfem::tuple
+ */
+template <typename T>
+struct is_tuple_of_tuples : std::false_type
+{
+};
+
+/**
+ * @brief Trait for checking if a type if a @p mfem::tuple containing only @p mfem::tuple
+ */
+template <typename... T>
+struct is_tuple_of_tuples<mfem::tuple<T...>>
+{
+   static constexpr bool value = (is_tuple<T>::value &&
+                                  ...);  ///< true/false result of type check
+};
+
+}  // namespace mfem

examples/dfem_advection_supg.cpp

@@ -0,0 +1,536 @@
+#include "dfem/dfem_refactor.hpp"
+#include "linalg/ode.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+class TimeStepEstimateQFunction
+{
+public:
+   TimeStepEstimateQFunction() = default;
+
+   using vecd = tensor<real_t, 2>;
+   using matd = tensor<real_t, 2, 2>;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const matd &dvdxi,
+      const real_t &rho0,
+      const matd &J0,
+      const matd &J,
+      const real_t &gamma,
+      const real_t &E,
+      const real_t &h0,
+      const real_t &order_v,
+      const real_t &w)
+   {
+      real_t dt_est = 0.0;
+      return mfem::tuple{dt_est};
+   }
+};
+
+void velocity(const Vector &c, Vector &u)
+{
+   const double x = c(0);
+   const double y = c(1);
+
+   u(0) = 0.5 - y;
+   u(1) = x - 0.5;
+}
+
+template <int problem = 0>
+real_t three_bodies_ic(const Vector &X)
+{
+   const real_t x = X(0);
+   const real_t y = X(1);
+   const real_t r0 = 0.15;
+   real_t x0 = 0.0;
+   real_t y0 = 0.0;
+
+   auto region = [&r0](const real_t x, const real_t y, const real_t x0,
+                       const real_t y0)
+   {
+      return sqrt(pow(x-x0, 2.0) + pow(y-y0, 2.0));
+   };
+
+   x0 = 0.25;
+   y0 = 0.5;
+   const real_t hump = 0.25 + 0.25 * cos(M_PI * region(x, y, x0, y0) / r0);
+   if (region(x, y, x0, y0) <= r0)
+   {
+      return hump;
+   }
+
+   if constexpr (problem == 1)
+   {
+      return 0.0;
+   }
+
+   x0 = 0.5;
+   y0 = 0.25;
+   const real_t cone = (1.0 - region(x, y, x0, y0) / r0);
+   if (region(x, y, x0, y0) <= r0)
+   {
+      return cone;
+   }
+
+   x0 = 0.5;
+   y0 = 0.75;
+   if ((region(x, y, x0, y0) <= r0) && (fabs(x - 0.5) >= 0.025 || y >= 0.85))
+   {
+      return 1.0;
+   }
+
+   return 0.0;
+}
+
+MFEM_HOST_DEVICE
+template <int dim = 2>
+tensor<real_t, dim> get_velocity(const tensor<real_t, dim>& x)
+{
+   return {0.5 - x(1), x(0) - 0.5};
+}
+
+MFEM_HOST_DEVICE
+template <int dim = 2>
+real_t compute_tau(
+   const tensor<real_t, dim>& b,
+   const tensor<real_t, dim, dim>& J,
+   const real_t& dt,
+   const int& p)
+{
+   const real_t h_min = calcsv(J, dim-1) / static_cast<real_t>(p);
+   real_t velocity_norm = sqrt(dot(b, b));
+   if (velocity_norm < 1e-12) { velocity_norm = 1e-12; }
+   auto tau_ugn_1 = h_min;
+   auto tau_ugn_2 = dt / 2.0;
+   auto tau = 1.0 / sqrt(1.0/pow(tau_ugn_1, 2) + 1.0/pow(tau_ugn_2, 2));
+   return tau;
+}
+
+template <int dim = 2>
+class SUPGMassQFunction
+{
+public:
+   SUPGMassQFunction(const real_t &dt, const int &p) :
+      dt(dt),
+      p(p)
+   {}
+
+   MFEM_HOST_DEVICE inline
+   auto operator() (
+      const real_t &k,
+      const tensor<real_t, dim>& x,
+      const tensor<real_t, dim, dim>& J,
+      const real_t& w) const
+   {
+      auto b = get_velocity(x);
+      auto tau = compute_tau(b, J, dt, p);
+
+      return mfem::tuple{tau * k * b * transpose(inv(J)) * det(J) * w};
+   }
+
+   const real_t &dt;  // time step
+   const int p;  // polynomial order
+};
+
+template <int dim = 2>
+class AdvQFunction
+{
+public:
+   AdvQFunction(const real_t &dt, const int &p, bool use_stabilization) :
+      dt(dt),
+      p(p),
+      use_stabilization(use_stabilization)
+   {};
+
+   MFEM_HOST_DEVICE inline
+   auto operator() (
+      const real_t &u,
+      const tensor<real_t, dim>& dudxi,
+      const tensor<real_t, dim>& x,
+      const tensor<real_t, dim, dim>& J,
+      const real_t& w) const
+   {
+      auto invJ = inv(J);
+      auto b = get_velocity(x);
+
+      // Advection
+      auto advection = -b * u;
+
+      if (use_stabilization)
+      {
+         auto residual = dot(b, (transpose(inv(J)) * dudxi));
+         auto tau = compute_tau(b, J, dt, p);
+         auto stab = tau * residual * b;
+         return mfem::tuple{(advection + stab) * transpose(invJ) * det(J) * w};
+      }
+      else
+      {
+         return mfem::tuple{advection * transpose(invJ) * det(J) * w};
+      }
+   }
+
+   const bool use_stabilization;
+   const real_t &dt;  // time step
+   const int p;  // polynomial order
+};
+
+template <int dim = 2>
+class AdvOp : public TimeDependentOperator
+{
+   static constexpr int Concentration = 0;
+   static constexpr int Coordinates = 1;
+
+   class AdvGradientOp : public Operator
+   {
+   public:
+      AdvGradientOp(const AdvOp &a, const Vector &x, real_t h) :
+         Operator(a.Height()),
+         a(a),
+         concentration_l(a.fes.GetVSize()),
+         h(h)
+      {
+         ParGridFunction g(&a.fes, concentration_l);
+         a.fes.GetProlongationMatrix()->Mult(x, g);
+         dRdu = a.adv->GetDerivative(Concentration, {&g}, {a.mesh_nodes});
+      }
+
+      void Mult(const Vector &k, Vector &y) const override
+      {
+         // column elimination for essential dofs
+         k_elim = k;
+         k_elim.SetSubVector(a.ess_tdof_list, 0.0);
+
+         dRdu->Mult(k_elim, y);
+         y *= h;
+         a.M->AddMult(k_elim, y);
+         if (a.use_stabilization)
+         {
+            a.Msupg_dk->AddMult(k_elim, y);
+         }
+
+         for (int i = 0; i < a.ess_tdof_list.Size(); i++)
+         {
+            y[a.ess_tdof_list[i]] = k[a.ess_tdof_list[i]];
+         }
+      }
+
+      const AdvOp &a;
+      mutable Vector concentration_l;
+      mutable Vector k_elim;
+      real_t h;
+      std::shared_ptr<DerivativeOperator> dRdu;
+   };
+
+   class AdvResidualOp : public Operator
+   {
+   public:
+      AdvResidualOp(const AdvOp &a, real_t dt, const Vector &x) :
+         Operator(a.Height()),
+         dt(dt),
+         a(a),
+         x(x),
+         u(x.Size()),
+         z(x.Size())
+      {}
+
+      void Mult(const Vector &k, Vector &R) const override
+      {
+         u = k;
+         u *= dt;
+         u += x;
+
+         a.M->Mult(k, R);
+         if (a.use_stabilization)
+         {
+            a.Msupg_dk->AddMult(k, R);
+         }
+         a.adv->AddMult(u, R);
+
+         R.SetSubVector(a.ess_tdof_list, 0.0);
+      }
+
+      Operator& GetGradient(const Vector &k) const override
+      {
+         u = k;
+         u *= dt;
+         u += x;
+
+         jacobian.reset(new AdvGradientOp(a, u, dt));
+         return *jacobian;
+
+         // fd_jacobian.reset(new FDJacobian(*this, k));
+         // return *fd_jacobian;
+      }
+
+      const AdvOp &a;
+      double dt;
+      Vector x;
+      mutable Vector u, z;
+
+      mutable std::shared_ptr<FDJacobian> fd_jacobian;
+
+      // AD Jacobian operator dRdu
+      mutable std::shared_ptr<AdvGradientOp> jacobian;
+   };
+
+public:
+   AdvOp(ParFiniteElementSpace &fes, const IntegrationRule &ir,
+         const Array<int> ess_tdof_list, const double &dt, const int &polynomial_order,
+         bool disable_tensor_product_structure = false, bool use_stabilization = false) :
+      TimeDependentOperator(fes.GetTrueVSize()),
+      ess_tdof_list(ess_tdof_list),
+      fes(fes),
+      Mform(&fes),
+      use_stabilization(use_stabilization)
+   {
+      auto mesh = fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>
+                   (mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      auto solutions = std::vector{FieldDescriptor{Concentration, &fes}};
+      auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes}};
+
+      adv = std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+      adv->DisableTensorProductStructure(disable_tensor_product_structure);
+      auto derivatives = std::integer_sequence<size_t, Concentration> {};
+      AdvQFunction<2> adv_qf(dt, polynomial_order, use_stabilization);
+      {
+         auto input_operators = mfem::tuple
+         {
+            Value<Concentration>{},
+            Gradient<Concentration>{},
+            Value<Coordinates>{},
+            Gradient<Coordinates>{},
+            Weight{}
+         };
+         auto output_operator = mfem::tuple{Gradient<Concentration>{}};
+
+         adv->AddDomainIntegrator(adv_qf, input_operators, output_operator, ir,
+                                  derivatives);
+      }
+      adv->SetParameters({mesh_nodes});
+
+      supg_mass = std::make_shared<DifferentiableOperator>(solutions, parameters,
+                                                           *mesh);
+      supg_mass->DisableTensorProductStructure(disable_tensor_product_structure);
+      SUPGMassQFunction<2> supg_mass_qf(dt, polynomial_order);
+      {
+         auto input_operators = mfem::tuple
+         {
+            Value<Concentration>{},
+            Value<Coordinates>{},
+            Gradient<Coordinates>{},
+            Weight{}
+         };
+         auto output_operator = mfem::tuple{Gradient<Concentration>{}};
+
+         supg_mass->AddDomainIntegrator(supg_mass_qf, input_operators, output_operator,
+                                        ir, derivatives);
+      }
+      supg_mass->SetParameters({mesh_nodes});
+
+      // Compute SUPG mass matrix by linearizing around a dummy variable
+      {
+         ParGridFunction g(&fes);
+         g = 1.0;
+         Msupg_dk = supg_mass->GetDerivative(Concentration, {&g}, {mesh_nodes});
+      }
+
+      Mform.AddDomainIntegrator(new MassIntegrator);
+      Mform.Assemble();
+      Mform.FormSystemMatrix(Array<int> {}, M);
+   }
+
+   void ImplicitSolve(const double dt, const Vector &x, Vector &k) override
+   {
+      auto residual = AdvResidualOp(*this, dt, x);
+
+      GMRESSolver krylov(MPI_COMM_WORLD);
+      krylov.SetRelTol(1e-6);
+      krylov.SetMaxIter(1000);
+      // krylov.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+      NewtonSolver newton(MPI_COMM_WORLD);
+      newton.SetOperator(residual);
+      newton.SetSolver(krylov);
+      newton.SetRelTol(1e-12);
+      newton.SetMaxIter(10);
+      // newton.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+      Vector zero;
+      k = x;
+      k.SetSubVector(ess_tdof_list, 0.0);
+      newton.Mult(zero, k);
+   }
+
+private:
+   std::shared_ptr<DifferentiableOperator> adv;
+   std::shared_ptr<DifferentiableOperator> supg_mass;
+   std::shared_ptr<DerivativeOperator> Msupg_dk;
+   Array<int> ess_tdof_list;
+   ParFiniteElementSpace &fes;
+   ParGridFunction *mesh_nodes = nullptr;
+   ParBilinearForm Mform;
+   OperatorHandle M;
+   const bool use_stabilization;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "../data/ref-square.mesh";
+   int polynomial_order = 1;
+   int ir_order = 2;
+   int refinements = 0;
+   real_t dt = 1.0;
+   real_t t_final = 2.0 * M_PI;
+   int vis_steps = 5;
+   bool disable_tensor_product_structure = false;
+   bool use_stabilization = false;
+   int problem = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&t_final, "-tf", "--t-final", "Final time; start time is 0.");
+   args.AddOption(&dt, "-dt", "--time-step", "Time step.");
+   args.AddOption(&vis_steps, "-vs", "--visualization-steps",
+                  "Visualize every n-th timestep.");
+   args.AddOption(&disable_tensor_product_structure, "-disable-tp", "--disable-tp",
+                  "-enable-tp", "--enable-tp", "");
+   args.AddOption(&use_stabilization, "-enable-stab", "--enable-stab",
+                  "-disable-stab", "--disable-stab", "");
+   args.AddOption(&problem, "-prob", "--problem", "problem number");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const int global_tdof = h1fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "#dofs " << global_tdof << "\n";
+   }
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   ParGridFunction concentration(&h1fes);
+
+   FunctionCoefficient *concentration_ic_coef = nullptr;
+   if (problem == 0)
+   {
+      concentration_ic_coef = new FunctionCoefficient(three_bodies_ic<0>);
+   }
+   else
+   {
+      concentration_ic_coef = new FunctionCoefficient(three_bodies_ic<1>);
+   }
+
+   Array<int> ess_tdof_list;
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+   AdvOp advdiff(h1fes, ir, ess_tdof_list, dt, polynomial_order,
+                 disable_tensor_product_structure, use_stabilization);
+
+   ODESolver *ode_solver = new SDIRK33Solver;
+   ode_solver->Init(advdiff);
+
+   Vector zero, x(h1fes.GetTrueVSize());
+   concentration.ProjectCoefficient(*concentration_ic_coef);
+   concentration.GetTrueDofs(x);
+
+   // print_vector(x);
+
+   ParGridFunction sol(&h1fes), err(&h1fes);
+   sol.SetFromTrueDofs(x);
+
+   err.ProjectCoefficient(*concentration_ic_coef);
+
+   real_t t = 0.0;
+
+   ParaViewDataCollection dc("dfem_advection_supg", &mesh);
+   dc.SetHighOrderOutput(true);
+   dc.SetLevelsOfDetail(polynomial_order);
+   dc.RegisterField("concentration", &sol);
+   dc.RegisterField("err", &err);
+   dc.SetCycle(0);
+   dc.SetTime(t);
+   dc.Save();
+
+   bool done = false;
+   for (int ti = 0; !done;)
+   {
+      real_t dt_real = std::min(dt, t_final - t);
+      ode_solver->Step(x, t, dt_real);
+      // print_vector(x);
+      ti++;
+
+      done = (t >= t_final - 1e-8*dt);
+
+      if (done || ti % vis_steps == 0)
+      {
+         sol.SetFromTrueDofs(x);
+
+         err.ProjectCoefficient(*concentration_ic_coef);
+         for (int i = 0; i < err.Size(); i++)
+         {
+            err[i] = abs(sol[i] - err[i]);
+         }
+
+         dc.SetCycle(ti);
+         dc.SetTime(t);
+         dc.Save();
+
+         if (Mpi::Root())
+         {
+            std::cout << "time step: " << ti << ", time: " << t << std::endl;
+         }
+      }
+   }
+
+   const real_t l2err = sol.ComputeL2Error(*concentration_ic_coef);
+   const real_t maxerr = sol.ComputeMaxError(*concentration_ic_coef);
+   if (Mpi::Root())
+   {
+      std::cout << "|u - u_ic|_L2 = " << l2err << "\n";
+      std::cout << "|u - u_ic|_max = " << maxerr << "\n";
+   }
+
+   return 0;
+}

examples/dfem_cfd.cpp

@@ -0,0 +1,932 @@
+#include "dfem/dfem_refactor.hpp"
+#include "general/tic_toc.hpp"
+#include "linalg/hypre.hpp"
+#include "linalg/solvers.hpp"
+#include "miniapps/autodiff/tadvector.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+constexpr int DIMENSION = 2;
+
+enum ProblemType
+{
+   CFD_EX_TEST = 0,
+   DFG_CFD1_TEST = 1,
+   DFG_CFD2_TEST = 2,
+};
+
+float clamp(float x, float lowerlimit = 0.0f, float upperlimit = 1.0f)
+{
+   if (x < lowerlimit) { return lowerlimit; }
+   if (x > upperlimit) { return upperlimit; }
+   return x;
+}
+
+struct CFD_TEST_CTX
+{
+   const int L = 1;
+} cfd_test_ctx;
+
+template <int dim = 2>
+struct VelocityMassQFunction
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim> &v,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      return mfem::tuple{v * det(J) * w};
+   }
+};
+
+template <int dim = 2>
+struct PressureMassQFunction
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(const real_t &p,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      return mfem::tuple{p * det(J) * w};
+   }
+};
+
+template <int dim = 2>
+struct NavierStokesMomentumConvectiveQFunction
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim> &v,
+                   const tensor<real_t, dim, dim> &dvdxi,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      return mfem::tuple{dot(dvdxi * inv(J), v) * det(J) * w};
+   }
+};
+
+template <int dim = 2>
+struct NavierStokesMomentumViscousQFunction
+{
+   NavierStokesMomentumViscousQFunction(real_t &kinematic_viscosity) :
+      kinematic_viscosity(kinematic_viscosity) {};
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim, dim> &dvdxi,
+                   const real_t &p,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      static constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      auto invJ = inv(J);
+      auto dvdx = dvdxi * invJ;
+      auto viscous_stress = -p * I + 2.0 * kinematic_viscosity * sym(dvdx);
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{(viscous_stress) * JxW};
+   }
+   const real_t kinematic_viscosity;
+};
+
+template <int dim = 2>
+struct NavierStokesContinuityQFunction
+{
+   NavierStokesContinuityQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim, dim> &dvdxi,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      return mfem::tuple{tr(dvdxi * inv(J)) * det(J) * w};
+   }
+};
+
+class ALEFSIOperator : public TimeDependentOperator
+{
+   static constexpr int Position = 0;
+   static constexpr int Velocity = 1;
+   static constexpr int Pressure = 2;
+
+   class ALEFSIResidual : public Operator
+   {
+   public:
+      class ALEFSIResJac : public Operator
+      {
+      public:
+         ALEFSIResJac(const ALEFSIResidual &res, const Vector &u) :
+            Operator(u.Size()),
+            res(res),
+            u(u)
+         {
+            fd_jacobian = std::make_shared<FDJacobian>(res, u);
+         }
+
+         void Mult(const Vector &x, Vector &y) const override
+         {
+            fd_jacobian->Mult(x, y);
+         }
+
+         const ALEFSIResidual &res;
+         const Vector u;
+         mutable std::shared_ptr<FDJacobian> fd_jacobian;
+      };
+
+      class ALEFSIJacPrec : public Solver
+      {
+      public:
+         ALEFSIJacPrec() : Solver() {}
+
+         void SetOperator(const Operator &jac) override
+         {
+            this->height = jac.Height();
+            this->width = jac.Width();
+
+            auto alefsi_jac = dynamic_cast<const ALEFSIResJac*>(&jac);
+            MFEM_VERIFY(alefsi_jac != nullptr, "invalid operator");
+
+            const ALEFSIResidual &res = alefsi_jac->res;
+            ALEFSIOperator &op = res.op;
+
+            Vector uv, up;
+            auto uptr = const_cast<Vector*>(&alefsi_jac->u);
+            uv.MakeRef(*uptr, 0, op.H1vtsize);
+            up.MakeRef(*uptr, op.H1vtsize, op.H1tsize);
+
+            auto x_gf = static_cast<ParGridFunction*>(op.H1vfes.GetParMesh()->GetNodes());
+            op.v_gf.SetFromTrueDofs(uv);
+            op.p_gf.SetFromTrueDofs(up);
+
+            HypreParMatrix Mv, Mp, Aconv, Avisc, B;
+
+            auto dMDv = op.fluid_velocity_mass->GetDerivative(Velocity, {&op.v_gf}, {x_gf});
+            auto dMDp = op.pressure_mass->GetDerivative(Pressure, {&op.p_gf}, {x_gf});
+            auto dFcvDv = op.fluid_momentum_convective->GetDerivative(Velocity, {&op.v_gf}, {x_gf});
+            auto dFvvDv = op.fluid_momentum_viscous->GetDerivative(Velocity, {&op.v_gf}, {&op.p_gf, x_gf});
+            auto dCDv = op.fluid_continuity->GetDerivative(Velocity, {&op.v_gf}, {&op.p_gf, x_gf});
+
+            dMDv->Assemble(Mv);
+            dMDp->Assemble(Mp);
+            dFcvDv->Assemble(Aconv);
+            dFvvDv->Assemble(Avisc);
+            dCDv->Assemble(B);
+
+            std::shared_ptr<HypreParMatrix> A0, A;
+            A0.reset(Add(1.0/res.gamma, Mv, 1.0, Aconv));
+            A.reset(Add(1.0, *A0, 1.0, Avisc));
+            auto Bt = B.Transpose();
+
+            Array2D<const HypreParMatrix*> blocks(2, 2);
+            blocks(0, 0) = A.get();
+            blocks(0, 1) = Bt;
+            blocks(1, 0) = &B;
+            blocks(1, 1) = &Mp;
+
+            Array2D<real_t> blockCoeff(2, 2);
+            blockCoeff(0, 0) = 1.0;
+            blockCoeff(0, 1) = -1.0;
+            blockCoeff(1, 0) = -1.0;
+            blockCoeff(1, 1) = 0.0;
+
+            K.reset(HypreParMatrixFromBlocks(blocks, &blockCoeff));
+
+            // std::ofstream kmout("K.m");
+            // kmout.precision(16);
+            // K->PrintMatlab(kmout);
+            // kmout.close();
+
+            Array<int> combined_ess_tdof(op.vel_ess_tdof.Size() + op.pres_ess_tdof.Size());
+            for (int i = 0; i < op.vel_ess_tdof.Size(); i++)
+            {
+               combined_ess_tdof[i] = op.vel_ess_tdof[i];
+            }
+            for (int i = 0; i < op.pres_ess_tdof.Size(); i++)
+            {
+               combined_ess_tdof[i + op.vel_ess_tdof.Size()] =
+                  op.pres_ess_tdof[i] + op.H1vtsize;
+            }
+
+            auto Ke = K->EliminateRowsCols(combined_ess_tdof);
+            delete Ke;
+
+            // std::ofstream kemout("Kelim.m");
+            // kemout.precision(16);
+            // K->PrintMatlab(kemout);
+            // kemout.close();
+
+            // auto fd_jacobian = std::make_shared<FDJacobian>(res, alefsi_jac->u);
+            // std::ofstream fdout("Kfd.m");
+            // fdout.precision(16);
+            // fd_jacobian->PrintMatlab(fdout);
+            // fdout.close();
+
+            // exit(0);
+
+            slu = std::make_shared<SuperLUSolver>(MPI_COMM_WORLD);
+            slu->SetPrintStatistics(false);
+            A_SLU = std::make_shared<SuperLURowLocMatrix>(*K);
+            slu->SetOperator(*A_SLU);
+         }
+
+         void Mult(const Vector &x, Vector &y) const override
+         {
+            // GMRESSolver krylov(MPI_COMM_WORLD);
+            // krylov.SetRelTol(1e-8);
+            // krylov.SetMaxIter(1000);
+            // krylov.SetOperator(*K);
+            // krylov.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+            // krylov.Mult(x, y);
+
+            slu->Mult(x, y);
+
+            // y = x;
+         }
+
+         std::shared_ptr<HypreParMatrix> K;
+         std::shared_ptr<SuperLURowLocMatrix> A_SLU;
+         std::shared_ptr<SuperLUSolver> slu;
+      };
+
+      ALEFSIResidual(ALEFSIOperator &op, const real_t &gamma, const Vector &S,
+                     const Vector &prevS) :
+         Operator(op.offsets.Last()),
+         op(op),
+         gamma(gamma),
+         prevu(prevS),
+         z(S.Size()),
+         v_gf(&op.H1vfes),
+         p_gf(&op.H1fes),
+         H1vtsize(op.H1vfes.GetTrueVSize()),
+         H1tsize(op.H1fes.GetTrueVSize()) {}
+
+      void Mult(const Vector &u, Vector &R) const override
+      {
+         auto uptr = const_cast<Vector*>(&u);
+
+         Vector uv, up;
+         uv.MakeRef(*uptr, 0, H1vtsize);
+         up.MakeRef(*uptr, H1vtsize, H1tsize);
+
+         Vector prevuv, prevup;
+         prevuv.MakeRef(prevu, 0, H1vtsize);
+         prevup.MakeRef(prevu, H1vtsize, H1tsize);
+
+         Vector Zv, Zp;
+         Zv.MakeRef(z, 0, H1vtsize);
+         Zp.MakeRef(z, H1vtsize, H1tsize);
+
+         Vector Rv, Rp;
+         Rv.MakeRef(R, 0, H1vtsize);
+         Rp.MakeRef(R, H1vtsize, H1tsize);
+
+         Rv = 0.0;
+         Rp = 0.0;
+
+         auto x_gf = static_cast<ParGridFunction*>(op.H1vfes.GetParMesh()->GetNodes());
+
+         subtract(uv, prevuv, Zv);
+         op.fluid_velocity_mass->SetParameters({x_gf});
+         op.fluid_velocity_mass->AddMult(Zv, Rv, 1.0/gamma);
+
+         op.fluid_momentum_convective->SetParameters({x_gf});
+         op.fluid_momentum_convective->AddMult(uv, Rv);
+
+         p_gf.SetFromTrueDofs(up);
+         // p_gf *= 1.0 / op.theta;
+         op.fluid_momentum_viscous->SetParameters({&p_gf, x_gf});
+         op.fluid_momentum_viscous->AddMult(uv, Rv);
+
+         p_gf.SetFromTrueDofs(up);
+         op.fluid_continuity->SetParameters({&p_gf, x_gf});
+         op.fluid_continuity->AddMult(uv, Rp, -1.0);
+
+         Rv.SetSubVector(op.vel_ess_tdof, 0.0);
+         Rp.SetSubVector(op.pres_ess_tdof, 0.0);
+      }
+
+      Operator& GetGradient(const Vector &u) const override
+      {
+         jacobian.reset(new ALEFSIResJac(*this, u));
+         return *jacobian;
+      }
+
+      ALEFSIOperator &op;
+      const real_t gamma;
+
+      mutable ParGridFunction v_gf, p_gf;
+
+      const int H1vtsize;
+      const int H1tsize;
+
+      mutable Vector z, prevu;
+
+      mutable std::shared_ptr<ALEFSIResJac> jacobian;
+   };
+
+public:
+   ALEFSIOperator(
+      real_t theta,
+      real_t &kinematic_viscosity,
+      ParFiniteElementSpace &H1vfes,
+      ParFiniteElementSpace &H1fes,
+      Array<int> &offsets,
+      Array<int> &fluid_domain_attr,
+      Array<int> &vel_ess_bdr,
+      Array<int> &pres_ess_bdr,
+      VectorCoefficient &vel_bdr_coeff,
+      Coefficient &pres_bdr_coeff,
+      const IntegrationRule &ir,
+      bool enable_tps) :
+      TimeDependentOperator(offsets.Last()),
+      theta(theta),
+      kinematic_viscosity(kinematic_viscosity),
+      offsets(offsets),
+      H1vfes(H1vfes),
+      H1fes(H1fes),
+      H1vtsize(H1vfes.GetTrueVSize()),
+      H1tsize(H1fes.GetTrueVSize()),
+      vel_ess_bdr(vel_ess_bdr),
+      pres_ess_bdr(pres_ess_bdr),
+      vel_bdr_coeff(&vel_bdr_coeff),
+      pres_bdr_coeff(&pres_bdr_coeff),
+      ir(ir),
+      prevS(offsets.Last()),
+      v_gf(&H1vfes),
+      p_gf(&H1fes)
+   {
+      auto mesh = H1vfes.GetParMesh();
+      x_gf = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *x_gf->ParFESpace();
+
+      H1vfes.GetEssentialTrueDofs(vel_ess_bdr, vel_ess_tdof);
+      H1fes.GetEssentialTrueDofs(pres_ess_bdr, pres_ess_tdof);
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Velocity, &H1vfes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Position, &mesh_fes}
+         };
+
+         fluid_velocity_mass =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         if (!enable_tps)
+         {
+            fluid_velocity_mass->DisableTensorProductStructure();
+         }
+
+         mfem::tuple inputs{Value<Velocity>{}, Gradient<Position>{}, Weight{}};
+         mfem::tuple outputs{Value<Velocity>{}};
+
+         auto mass_qf = VelocityMassQFunction<DIMENSION> {};
+         auto derivatives = std::integer_sequence<size_t, Velocity> {};
+         fluid_velocity_mass->AddDomainIntegrator(mass_qf, inputs, outputs, ir,
+                                                  fluid_domain_attr, derivatives);
+      }
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Velocity, &H1vfes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Position, &mesh_fes}
+         };
+
+         fluid_momentum_convective =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         if (!enable_tps)
+         {
+            fluid_momentum_convective->DisableTensorProductStructure();
+         }
+
+         mfem::tuple inputs
+         {
+            Value<Velocity>{},
+            Gradient<Velocity>{},
+            Gradient<Position>{},
+            Weight{}
+         };
+
+         mfem::tuple outputs{Value<Velocity>{}};
+
+         auto momentum_qf = NavierStokesMomentumConvectiveQFunction<DIMENSION> {};
+         auto derivatives = std::integer_sequence<size_t, Velocity> {};
+         fluid_momentum_convective->AddDomainIntegrator(
+            momentum_qf, inputs, outputs, ir, fluid_domain_attr, derivatives);
+      }
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Velocity, &H1vfes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Pressure, &H1fes},
+            FieldDescriptor{Position, &mesh_fes}
+         };
+
+         fluid_momentum_viscous =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         if (!enable_tps)
+         {
+            fluid_momentum_viscous->DisableTensorProductStructure();
+         }
+
+         mfem::tuple inputs
+         {
+            Gradient<Velocity>{},
+            Value<Pressure>{},
+            Gradient<Position>{},
+            Weight{}
+         };
+
+         mfem::tuple outputs{Gradient<Velocity>{}};
+
+         auto momentum_qf =
+            NavierStokesMomentumViscousQFunction<DIMENSION>(kinematic_viscosity);
+         auto derivatives = std::integer_sequence<size_t, Velocity> {};
+         fluid_momentum_viscous->AddDomainIntegrator(
+            momentum_qf, inputs, outputs, ir, fluid_domain_attr, derivatives);
+      }
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Velocity, &H1vfes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Pressure, &H1fes},
+            FieldDescriptor{Position, &mesh_fes}
+         };
+
+         fluid_continuity =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         if (!enable_tps)
+         {
+            fluid_continuity->DisableTensorProductStructure();
+         }
+
+         mfem::tuple inputs{Gradient<Velocity>{}, Gradient<Position>{}, Weight{}};
+         mfem::tuple outputs{Value<Pressure>{}};
+         auto continuity_qf = NavierStokesContinuityQFunction<DIMENSION> {};
+         auto derivatives = std::integer_sequence<size_t, Velocity> {};
+         fluid_continuity->AddDomainIntegrator(
+            continuity_qf, inputs, outputs, ir, fluid_domain_attr, derivatives);
+      }
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Pressure, &H1fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Position, &mesh_fes}
+         };
+
+         pressure_mass =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         if (!enable_tps)
+         {
+            pressure_mass->DisableTensorProductStructure();
+         }
+
+         mfem::tuple inputs{Value<Pressure>{}, Gradient<Position>{}, Weight{}};
+         mfem::tuple outputs{Value<Pressure>{}};
+         auto pressure_mass_qf = PressureMassQFunction<DIMENSION> {};
+         auto derivatives = std::integer_sequence<size_t, Pressure> {};
+         pressure_mass->AddDomainIntegrator(
+            pressure_mass_qf, inputs, outputs, ir, fluid_domain_attr, derivatives);
+      }
+   }
+
+   void SetTime(const real_t t) override
+   {
+      vel_bdr_coeff->SetTime(t);
+      pres_bdr_coeff->SetTime(t);
+   }
+
+   void Step(Vector &S, real_t &t, const real_t &dt)
+   {
+      this->SetTime(t);
+      prevS = S;
+
+      this->SetTime(t + dt);
+      Vector Sv, Sp;
+      Sv.MakeRef(S, 0, H1vtsize);
+      Sp.MakeRef(S, H1vtsize, H1tsize);
+
+      v_gf.SetFromTrueDofs(Sv);
+      v_gf.ProjectBdrCoefficient(*vel_bdr_coeff, vel_ess_bdr);
+      v_gf.GetTrueDofs(Sv);
+
+      p_gf.SetFromTrueDofs(Sp);
+      p_gf.ProjectBdrCoefficient(*pres_bdr_coeff, pres_ess_bdr);
+      p_gf.GetTrueDofs(Sp);
+
+      ALEFSIResidual residual(*this, dt, S, prevS);
+
+      ALEFSIOperator::ALEFSIResidual::ALEFSIJacPrec prec;
+
+      GMRESSolver krylov(MPI_COMM_WORLD);
+      krylov.SetRelTol(1e-4);
+      krylov.SetMaxIter(1000);
+      krylov.SetKDim(300);
+      // krylov.SetPreconditioner(prec);
+      krylov.SetPrintLevel(IterativeSolver::PrintLevel().FirstAndLast());
+
+      NewtonSolver newton(MPI_COMM_WORLD);
+      newton.SetOperator(residual);
+      newton.SetSolver(krylov);
+      newton.SetRelTol(1e-6);
+      newton.SetMaxIter(10);
+      newton.SetPrintLevel(IterativeSolver::PrintLevel().Iterations());
+
+      Vector zero;
+      newton.Mult(zero, S);
+
+      t += dt;
+   }
+
+   real_t theta;
+   real_t kinematic_viscosity;
+   Array<int> offsets;
+   std::shared_ptr<DifferentiableOperator> fluid_momentum_convective;
+   std::shared_ptr<DifferentiableOperator> fluid_momentum_viscous;
+   std::shared_ptr<DifferentiableOperator> fluid_continuity;
+   std::shared_ptr<DifferentiableOperator> fluid_velocity_mass;
+   std::shared_ptr<DifferentiableOperator> pressure_mass;
+
+   ParGridFunction *x_gf, v_gf, p_gf;
+
+   Vector prevS;
+
+   const Array<int> vel_ess_bdr, pres_ess_bdr;
+   Array<int> vel_ess_tdof, pres_ess_tdof;
+
+   VectorCoefficient *vel_bdr_coeff;
+   Coefficient *pres_bdr_coeff;
+
+   ParFiniteElementSpace &H1vfes;
+   ParFiniteElementSpace &H1fes;
+   const int H1vtsize, H1tsize;
+   IntegrationRule ir;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+   Hypre::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "";
+   int polynomial_order_velocity = 2;
+   int refinements = 0;
+   int problem_type = 0;
+   real_t t_final = 0.0;
+   real_t dt = 1e-3;
+   real_t kinematic_viscosity = 1.0;
+   int vis_steps = 1;
+   real_t theta = 1.0;
+   bool enable_tps = true;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&problem_type, "-prob", "--problem", "Problem #");
+   args.AddOption(&polynomial_order_velocity, "-ov", "--order-velocity", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&t_final, "-tf", "--tf", "");
+   args.AddOption(&dt, "-dt", "--dt", "");
+   args.AddOption(&kinematic_viscosity, "-kv", "--kinematic-viscosity", "");
+   args.AddOption(&theta, "-theta", "--theta", "");
+   args.AddOption(&vis_steps, "-vs", "--visualization-steps",
+                  "Visualize every n-th timestep.");
+   args.AddOption(&enable_tps, "-tps", "--enable-tps", "-no-tps",
+                  "--no-tps", "Enable tensor product structure for quad/hex.");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   int polynomial_order_pressure = polynomial_order_velocity - 1;
+
+   Mesh mesh_serial;
+   if (problem_type == ProblemType::CFD_EX_TEST)
+   {
+      mesh_serial = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL);
+      mesh_serial.EnsureNodes();
+      auto nodes = mesh_serial.GetNodes();
+      *nodes -= 0.5;
+      *nodes *= cfd_test_ctx.L;
+   }
+   else if (problem_type == ProblemType::DFG_CFD1_TEST ||
+            problem_type == ProblemType::DFG_CFD2_TEST)
+   {
+      mesh_serial = Mesh::LoadFromFile(mesh_file);
+      // Array<int> domains(1);
+      // domains[0] = 1;
+      // mesh_serial = SubMesh::CreateFromDomain(m, domains);
+      mesh_serial.EnsureNodes();
+   }
+   else
+   {
+      MFEM_ABORT("invalid problem type");
+   }
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.EnsureNodes();
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   H1_FECollection velocity_fec(polynomial_order_velocity, dim);
+   H1_FECollection pressure_fec(polynomial_order_pressure);
+
+   ParFiniteElementSpace H1vfes(&mesh, &velocity_fec, dim);
+   ParFiniteElementSpace H1fes(&mesh, &pressure_fec);
+
+   HYPRE_BigInt global_size_velocity = H1vfes.GlobalTrueVSize();
+   HYPRE_BigInt global_size_pressure = H1fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "Number of velocity unknowns: " << global_size_velocity << "\n";
+      out << "Number of pressure unknowns: " << global_size_pressure << "\n";
+   }
+
+   Array<int> fluid_domain_attr(mesh.attributes.Max());
+   fluid_domain_attr = 1;
+
+   const IntegrationRule &integration_rule =
+      IntRules.Get(H1vfes.GetFE(0)->GetGeomType(),
+                   2 * H1vfes.GetFE(0)->GetOrder() + 1);
+
+   Array<int> vel_ess_attr(mesh.bdr_attributes.Max());
+   if (problem_type == ProblemType::CFD_EX_TEST)
+   {
+      vel_ess_attr = 1;
+   }
+   else if (problem_type == ProblemType::DFG_CFD1_TEST ||
+            problem_type == ProblemType::DFG_CFD2_TEST)
+   {
+      // everywhere
+      vel_ess_attr = 1;
+      // beam
+      vel_ess_attr[4] = 0;
+      // outlet
+      vel_ess_attr[1] = 0;
+   }
+
+   Array<int> pres_ess_attr(mesh.bdr_attributes.Max());
+   if (problem_type == ProblemType::CFD_EX_TEST)
+   {
+      pres_ess_attr = 1;
+   }
+   else if (problem_type == ProblemType::DFG_CFD1_TEST ||
+            problem_type == ProblemType::DFG_CFD2_TEST)
+   {
+      // everywhere
+      pres_ess_attr = 0;
+      // outlet
+      // pres_ess_attr[1] = 1;
+   }
+
+   Array<int> block_offsets(3);
+   block_offsets[0] = 0;
+   block_offsets[1] = H1vfes.GetTrueVSize();
+   block_offsets[2] = H1fes.GetTrueVSize();
+   block_offsets.PartialSum();
+
+   BlockVector S(block_offsets, Device::GetDeviceMemoryType());
+
+   ParGridFunction v_gf(&H1vfes), p_gf(&H1fes);
+
+   std::function<void(const Vector &, real_t, Vector &)> velocity_exact;
+   if (problem_type == ProblemType::CFD_EX_TEST)
+   {
+      velocity_exact = [nu = kinematic_viscosity](
+                          const Vector &coords, real_t t, Vector &u)
+      {
+         const real_t x = coords(0);
+         const real_t y = coords(1);
+         const real_t f = exp(-4.0 * nu * M_PI * M_PI * t);
+         u(0) = -sin(2.0 * M_PI * y) * f;
+         u(1) = sin(2.0 * M_PI * x) * f;
+      };
+   }
+   else if (problem_type == ProblemType::DFG_CFD1_TEST ||
+            problem_type == ProblemType::DFG_CFD2_TEST)
+   {
+      velocity_exact = [problem_type](const Vector &coords, real_t t, Vector &u)
+      {
+         const real_t x = coords(0);
+         const real_t y = coords(1);
+         const real_t H = 0.41;
+         real_t U = 0.3;
+         if (problem_type == DFG_CFD2_TEST)
+         {
+            U = 1.5;
+         }
+         auto smoothstep = [t](const real_t edge0, const real_t edge1, real_t x)
+         {
+            x = clamp((x - edge0) / (edge1 - edge0));
+            return x * x * (3.0 - 2.0 * x);
+         };
+         if (x == 0.0)
+         {
+            u(0) = 4.0 * U * y * (H - y) / powf(H, 2.0) * smoothstep(0.0, 0.1, t);
+         }
+         else
+         {
+            u(0) = 0.0;
+         }
+         u(1) = 0.0;
+      };
+   }
+
+   std::function<real_t(const Vector &, real_t)> pressure_exact;
+   if (problem_type == ProblemType::CFD_EX_TEST)
+   {
+      pressure_exact = [nu = kinematic_viscosity](const Vector &coords, real_t t)
+      {
+         const real_t x = coords(0);
+         const real_t y = coords(1);
+         const real_t f = exp(-8.0 * nu * M_PI * M_PI * t);
+         return -cos(2.0 * M_PI * x) * cos(2.0 * M_PI * y) * f;
+      };
+   }
+   else if (problem_type == ProblemType::DFG_CFD1_TEST ||
+            problem_type == ProblemType::DFG_CFD2_TEST)
+   {
+      pressure_exact = [](const Vector &, real_t)
+      {
+         return 0.0;
+      };
+   }
+
+   VectorFunctionCoefficient vel_exact(dim, velocity_exact);
+   FunctionCoefficient pres_exact(pressure_exact);
+
+   v_gf.ProjectCoefficient(vel_exact);
+   v_gf.GetTrueDofs(S.GetBlock(0));
+
+   p_gf.ProjectCoefficient(pres_exact);
+   p_gf.GetTrueDofs(S.GetBlock(1));
+
+   ALEFSIOperator alefsi(
+      theta,
+      kinematic_viscosity,
+      H1vfes,
+      H1fes,
+      block_offsets,
+      fluid_domain_attr,
+      vel_ess_attr,
+      pres_ess_attr,
+      vel_exact,
+      pres_exact,
+      integration_rule,
+      enable_tps);
+
+   real_t t = 0.0;
+   out << "time step: " << dt << "\n";
+   real_t t_old;
+   bool last_step = false;
+
+   ParGridFunction verr_gf(&H1vfes), vex_gf(&H1vfes);
+   verr_gf = 0.0;
+
+   ParGridFunction perr_gf(&H1fes), pex_gf(&H1fes);
+   perr_gf = 0.0;
+
+   ParaViewDataCollection dc("dfem_cfd", &mesh);
+   dc.SetHighOrderOutput(true);
+   dc.SetLevelsOfDetail(polynomial_order_velocity);
+   dc.RegisterField("velocity", &v_gf);
+   dc.RegisterField("pressure", &p_gf);
+   if (problem_type == ProblemType::CFD_EX_TEST)
+   {
+      dc.RegisterField("velocity_exact", &vex_gf);
+      dc.RegisterField("velocity_error", &verr_gf);
+      dc.RegisterField("pressure_exact", &pex_gf);
+      dc.RegisterField("pressure_error", &perr_gf);
+   }
+   dc.SetCycle(0);
+   dc.SetTime(0);
+   dc.Save();
+
+   for (int ti = 1; !last_step; ti++)
+   {
+      if (t + dt >= t_final - 1e-8*dt)
+      {
+         dt = t_final - t;
+         last_step = true;
+      }
+
+      alefsi.Step(S, t, dt);
+
+      if (last_step || (ti % vis_steps == 0))
+      {
+         v_gf.SetFromTrueDofs(S.GetBlock(0));
+         p_gf.SetFromTrueDofs(S.GetBlock(1));
+
+         if (Mpi::Root())
+         {
+            out << "step " << std::setw(5) << ti
+                << ",\tt = " << std::setw(5) << std::setprecision(4) << t
+                << ",\tdt = " << std::setw(5) << std::setprecision(6) << dt;
+            out << std::endl;
+         }
+
+         dc.SetCycle(ti);
+         dc.SetTime(t);
+         dc.Save();
+      }
+   }
+
+   if (problem_type == ProblemType::CFD_EX_TEST)
+   {
+      vel_exact.SetTime(t);
+      vex_gf.ProjectCoefficient(vel_exact);
+      real_t vel_l2err = v_gf.ComputeL2Error(vel_exact);
+
+      pres_exact.SetTime(t);
+      pex_gf.ProjectCoefficient(pres_exact);
+      real_t pres_l2err = p_gf.ComputeL2Error(pres_exact);
+
+      if (Mpi::Root())
+      {
+         out << "|u - u_exact|_L2 = " << vel_l2err
+             << "\n|p - p_exact|_L2 = " << pres_l2err;
+      }
+
+      for (int i = 0; i < verr_gf.Size(); i++)
+      {
+         verr_gf(i) = abs(vex_gf(i) - v_gf(i));
+      }
+
+      for (int i = 0; i < perr_gf.Size(); i++)
+      {
+         perr_gf(i) = abs(pex_gf(i) - p_gf(i));
+      }
+
+      dc.Save();
+   }
+
+   if (problem_type == ProblemType::DFG_CFD1_TEST ||
+       problem_type == ProblemType::DFG_CFD2_TEST)
+   {
+      DenseMatrix points(dim, 1);
+      Vector pointA(2), pointB(2);
+
+      pointA(0) = 0.25;
+      pointA(1) = 0.2;
+
+      pointB(0) = 0.15;
+      pointB(1) = 0.2;
+
+      Array<int> elem_ids;
+      Array<IntegrationPoint> ips;
+      points.SetCol(0, pointA);
+      mesh.FindPoints(points, elem_ids, ips);
+      real_t pA = p_gf.GetValue(elem_ids[0], ips[0]);
+
+      points.SetCol(0, pointB);
+      mesh.FindPoints(points, elem_ids, ips);
+      real_t pB = p_gf.GetValue(elem_ids[0], ips[0]);
+
+      out << "p(B) - p(A) = " << pB - pA << "\n";
+   }
+
+   Hypre::Finalize();
+   Mpi::Finalize();
+
+   return 0;
+}

examples/dfem_dg.cpp

@@ -0,0 +1,123 @@
+#include "dfem/dfem_refactor.hpp"
+#include "fem/bilininteg.hpp"
+#include "fem/coefficient.hpp"
+#include "linalg/auxiliary.hpp"
+#include "linalg/hypre.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int main(int argc, char *argv[])
+{
+   Mpi::Init();
+   int num_procs = Mpi::WorldSize();
+   int myid = Mpi::WorldRank();
+   Hypre::Init();
+
+   const char *mesh_file = "../data/ref-square.mesh";
+   int polynomial_order = 1;
+   int ir_order = 2;
+   int refinements = 1;
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   L2_FECollection fec(polynomial_order, dim, BasisType::GaussLobatto);
+   ParFiniteElementSpace fes(&mesh, &fec);
+
+   const IntegrationRule &ir = IntRules.Get(fes.GetFE(0)->GetGeomType(),
+                                            ir_order * fec.GetOrder());
+   const IntegrationRule &ir_face = IntRules.Get(
+                                       fes.GetTraceElement(0, fes.GetMesh()->GetFaceGeometry(0))->GetGeomType(),
+                                       ir_order * fec.GetOrder());
+   ParGridFunction u(&fes);
+
+   // // -\nabla \cdot (\nabla u + p * I) -> (\nabla u + p * I, \nabla v)
+   // auto advection_kernel = [](const tensor<double, 2> &dudxi,
+   //                            const tensor<double, 2, 2> &J,
+   //                            const double &w)
+   // {
+   //    constexpr tensor<double, 2> b{1.0, 1.0};
+   //    return std::tuple{dot(b, dudxi * inv(J)) * det(J) * w};
+   // };
+
+   // std::tuple argument_operators_0{Gradient{"quantity"}, Gradient{"coordinates"}, Weight{"integration_weights"}};
+   // std::tuple output_operator_0{Value{"quantity"}};
+   // ElementOperator op_0{advection_kernel, argument_operators_0, output_operator_0};
+
+   // std::array solutions{FieldDescriptor{&fes, "quantity"}};
+   // std::array parameters{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   // DifferentiableOperator advection_op{solutions, parameters, std::tuple{op_0}, mesh, ir};
+
+   // auto adv_du = advection_op.template GetDerivativeWrt<0>({&u}, {mesh_nodes});
+   // HypreParMatrix A;
+   // adv_du->Assemble(A);
+
+   // std::ofstream mmatofs("dfem_mat.dat");
+   // A.PrintMatlab(mmatofs);
+   // mmatofs.close();
+
+   auto trace_kernel = [](const double &uL, const double &uR, const double &J,
+                          const double &w)
+   {
+      return std::tuple{1.0 / J * w};
+   };
+
+   std::tuple argument_operators_0
+   {
+      FaceValueLeft{"quantity"},
+      FaceValueRight{"quantity"},
+      Gradient{"coordinates"},
+      Weight{"integration_weights"}
+   };
+   std::tuple output_operator_0{Value{"quantity"}};
+   FaceElementOperator op_0{trace_kernel, argument_operators_0, output_operator_0};
+
+   std::array solutions{FieldDescriptor{&fes, "quantity"}};
+   std::array parameters{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator trace_op{solutions, parameters, std::tuple{op_0}, mesh, ir_face};
+
+   auto vector_func = [](const Vector &, Vector &u)
+   {
+      u = 1.0;
+   };
+
+   VectorFunctionCoefficient vel_coeff(dim, vector_func);
+
+   ParBilinearForm adv_form(&fes);
+   constexpr double alpha = 1.0;
+   auto integ = new ConvectionIntegrator(vel_coeff, alpha);
+   integ->SetIntRule(&ir);
+   adv_form.AddInteriorFaceIntegrator(
+      new NonconservativeDGTraceIntegrator(vel_coeff, alpha));
+   // adv_form.AddDomainIntegrator(integ);
+   adv_form.Assemble();
+   adv_form.Finalize();
+
+   auto K = adv_form.ParallelAssemble();
+   std::ofstream kmatofs("mfem_mat.dat");
+   K->PrintMatlab(kmatofs);
+   kmatofs.close();
+
+   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" << mesh << u << std::flush;
+
+   return 0;
+}

examples/dfem_elasticity.cpp

@@ -0,0 +1,150 @@
+#include "dfem.hpp"
+
+int main(int argc, char *argv[])
+{
+   Mpi::Init();
+
+   std::cout << std::setprecision(9);
+
+   const char *mesh_file = "../data/star.mesh";
+   int polynomial_order = 1;
+   int refinements = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.ParseCheck();
+
+   Mesh mesh_serial(mesh_file, 1, 1);
+   mesh_serial.SetCurvature(1);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   const int dim = mesh_serial.Dimension();
+
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+   mesh_serial.Clear();
+
+   constexpr int vdim = 2;
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   Array<int> ess_tdof_list;
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   std::cout << "nqpts = " << ir.GetNPoints() << std::endl;
+   std::cout << "ndofs = " << h1fes.GlobalTrueVSize() << std::endl;
+
+   ParGridFunction u(&h1fes);
+
+   auto exact_solution = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      u(0) = x*x + y;
+      u(1) = x + 0.5*y*y;
+   };
+
+   VectorFunctionCoefficient exact_solution_coeff(dim, exact_solution);
+
+   auto elasticity_kernel = [](tensor<double, 2, 2> &dudxi,
+                               tensor<double, 2, 2> &J,
+                               double &w)
+   {
+      using mfem::internal::tensor;
+      using mfem::internal::IsotropicIdentity;
+
+      double lambda, mu;
+      {
+         lambda = 1.0;
+         mu = 1.0;
+      }
+      static constexpr auto I = IsotropicIdentity<2>();
+      auto eps = sym(dudxi * inv(J));
+      auto JxW = transpose(inv(J)) * det(J) * w;
+      auto r = (lambda * tr(eps) * I + 2.0 * mu * eps) * JxW;
+      return r;
+   };
+
+   tensor<double, 2, 2> dudxi, s_dudxi, J;
+   double w = 1.0;
+
+   enzyme::get<0>
+   (enzyme::autodiff<enzyme::Forward,
+    enzyme::DuplicatedNoNeed<tensor<double, 2, 2>>>
+    (+elasticity_kernel,
+     enzyme::Duplicated<tensor<double, 2, 2> *>(&dudxi, &s_dudxi),
+     enzyme::Const<tensor<double, 2, 2>*>(&J),
+     enzyme::Const<double*>(&w)));
+
+   // std::tuple input_descriptors = {Gradient{"displacement"}, Gradient{"coordinates"}, Weight{"integration_weight"}};
+   // std::tuple output_descriptors = {Gradient{"displacement"}};
+   // ElementOperator qf {elasticity_kernel, input_descriptors, output_descriptors};
+
+   // ElementOperator forcing_qf
+   // {
+   //    [](tensor<double, 2> x, tensor<double, 2, 2> J, double w)
+   //    {
+   //       double lambda, mu;
+   //       {
+   //          lambda = 1.0;
+   //          mu = 1.0;
+   //       }
+   //       auto f = x;
+   //       f(0) = 4.0*mu + 2.0*lambda;
+   //       f(1) = 2.0*mu + lambda;
+   //       return f * det(J) * w;
+   //    },
+   //    // inputs
+   //    std::tuple{
+   //       Value{"coordinates"},
+   //       Gradient{"coordinates"},
+   //       Weight{"integration_weight"}},
+   //    // outputs
+   //    std::tuple{
+   //       Value{"displacement"}}
+   // };
+
+   // std::vector<Field> solutions{{&u, "displacement"}};
+   // std::vector<Field> parameters{{mesh.GetNodes(), "coordinates"}};
+   // std::vector<Field> dependent_fields{{&u, "displacement"}};
+   // DifferentiableForm dop(solutions, parameters, dependent_fields, mesh);
+
+   // dop.AddElementOperator<AD::Enzyme>(qf, ir);
+   // dop.AddElementOperator<AD::None>(forcing_qf, ir);
+   // dop.SetEssentialTrueDofs(ess_tdof_list);
+
+   // GMRESSolver gmres(MPI_COMM_WORLD);
+   // gmres.SetRelTol(1e-12);
+   // gmres.SetMaxIter(5000);
+   // gmres.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+   // NewtonSolver newton(MPI_COMM_WORLD);
+   // newton.SetSolver(gmres);
+   // newton.SetOperator(dop);
+   // newton.SetRelTol(1e-12);
+   // newton.SetMaxIter(100);
+   // newton.SetPrintLevel(1);
+
+   // u = 1e-6;
+   // u.ProjectBdrCoefficient(exact_solution_coeff, ess_bdr);
+   // Vector x;
+   // u.GetTrueDofs(x);
+
+   // Vector zero;
+   // newton.Mult(zero, x);
+
+   // u.Distribute(x);
+
+   // std::cout << "|u-u_ex|_L2 = " << u.ComputeL2Error(exact_solution_coeff) << "\n";
+
+   return 0;
+}

examples/dfem_elasticity_vjp.cpp

@@ -0,0 +1,458 @@
+#include "dfem/dfem_refactor.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+constexpr int DIMENSION = 2;
+
+template <typename Material, int dim = DIMENSION>
+struct StressQFunction
+{
+   StressQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudxi,
+      const tensor<real_t, dim, dim> &J,
+      const double &w,
+      const double &E) const
+   {
+      auto invJ = inv(J);
+      auto dudX = dudxi * invJ;
+      auto P = material(dudX, E);
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{P * JxW};
+   }
+
+   Material material;
+};
+
+template <int dim = DIMENSION>
+struct ParameterizedNeoHookean
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudX, real_t E) const
+   {
+      constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      const real_t mu = 0.5 * E / (1.0 + nu);
+      const real_t lambda = 2.0 * mu * nu / (1.0 - 2.0 * nu);
+
+      auto F = dudX + I;
+      auto J = det(F);
+      auto logJ = log(J);
+      return mu*F + (lambda*logJ - mu)*inv(transpose(F));
+   }
+
+   double nu;
+};
+
+/**
+ * Strain energy density
+ */
+template <int dim = DIMENSION>
+struct ParameterizedNeoHookeanEnergy
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+    const tensor<real_t, dim, dim> &dudxi,
+    const tensor<real_t, dim, dim> &J,
+    const double &w,
+    const double &E) const
+   {
+      constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      const real_t mu = 0.5 * E / (1.0 + nu);
+      const real_t lambda = 2.0 * mu * nu / (1.0 - 2.0 * nu);
+
+      auto invJ = inv(J);
+      auto dudX = dot(dudxi, invJ);
+      auto tr_C_minus_I = 2 * tr(dudX) + inner(dudX, dudX);
+      auto logdetF = log(det(dudX + I));
+      auto psi = 0.5*mu*tr_C_minus_I - mu*logdetF + 0.5*lambda*logdetF*logdetF;
+      auto dV = det(J) * w;
+      return mfem::tuple{psi*dV};
+   }
+
+   double nu;
+};
+
+
+class ElasticityOperator : public Operator
+{
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+   static constexpr int ElasticModulus = 2;
+
+public:
+   class ElasticityJacobianPreconditioner : public Solver
+   {
+   public:
+      ElasticityJacobianPreconditioner() : Solver() {}
+
+      void SetOperator(const Operator &op) override
+      {
+         this->height = op.Height();
+         this->width = op.Width();
+
+         auto elasticity_jacobian = dynamic_cast<const ElasticityJacobianOperator*>(&op);
+         MFEM_VERIFY(elasticity_jacobian != nullptr, "invalid operator");
+
+         A = std::make_shared<HypreParMatrix>();
+         elasticity_jacobian->momentum_du->Assemble(*A);
+         auto Ae = A->EliminateRowsCols(
+                      elasticity_jacobian->elasticity->displacement_ess_tdof);
+         delete Ae;
+
+         amg = std::make_shared<HypreBoomerAMG>();
+         amg->SetOperator(*A);
+         amg->SetPrintLevel(0);
+         amg->SetSystemsOptions(
+            elasticity_jacobian->elasticity->mesh_nodes->ParFESpace()->GetMesh()->Dimension(),
+            true);
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         amg->Mult(x, y);
+      }
+
+      std::shared_ptr<HypreParMatrix> A;
+      std::shared_ptr<HypreBoomerAMG> amg;
+   };
+
+   class ElasticityJacobianOperator : public Operator
+   {
+   public:
+      ElasticityJacobianOperator(const ElasticityOperator *elasticity,
+                                 const Vector &x) :
+         Operator(elasticity->Height()),
+         elasticity(elasticity),
+         z(elasticity->Height())
+      {
+         ParGridFunction u(&elasticity->displacement_fes);
+         u.SetFromTrueDofs(x);
+
+         auto mesh_nodes = static_cast<ParGridFunction*>
+                           (elasticity->displacement_fes.GetParMesh()->GetNodes());
+         momentum_du = elasticity->momentum->GetDerivative(Displacement, {&u}, {mesh_nodes, elasticity->modulus});
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         z = x;
+         z.SetSubVector(elasticity->displacement_ess_tdof, 0.0);
+
+         momentum_du->Mult(z, y);
+
+         for (int i = 0; i < elasticity->displacement_ess_tdof.Size(); i++)
+         {
+            y[elasticity->displacement_ess_tdof[i]] =
+               x[elasticity->displacement_ess_tdof[i]];
+         }
+      }
+
+      const ElasticityOperator *elasticity;
+      std::shared_ptr<DerivativeOperator> momentum_du;
+      mutable Vector z;
+   };
+
+   ElasticityOperator(ParFiniteElementSpace &displacement_fes,
+                      Array<int> &vel_ess_tdofs,
+                      const IntegrationRule &displacement_ir,
+                      ParGridFunction& elastic_modulus) :
+      Operator(displacement_fes.GetTrueVSize()),
+      density(1.0e3),
+      displacement_ess_tdof(vel_ess_tdofs),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      body_force(displacement_fes.GetTrueVSize()),
+      modulus_fes(*elastic_modulus.ParFESpace()),
+      modulus(&elastic_modulus)
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Displacement, &displacement_fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Coordinates, &mesh_fes},
+            FieldDescriptor{ElasticModulus, &modulus_fes}
+         };
+
+         momentum =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+         momentum->DisableTensorProductStructure();
+
+         mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, Weight{}, Value<ElasticModulus>{}};
+         mfem::tuple outputs{Gradient<Displacement>{}};
+
+         using Material = ParameterizedNeoHookean<DIMENSION>;
+         auto material = Material{.nu = 0.25};
+         auto qfunction = StressQFunction<Material, DIMENSION> {.material = material};
+         auto derivatives = std::integer_sequence<size_t, Displacement> {};
+         Array<int> solid_domain_attr(mesh->attributes.Max());
+         solid_domain_attr[0] = 1;
+         momentum->AddDomainIntegrator(
+            qfunction, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+         
+         momentum->SetParameters({mesh_nodes, &elastic_modulus});
+      }
+
+      {
+         Vector g(DIMENSION);
+         g = 0.0;
+         g(1) = -2.0/3.0 * 1e-3;
+
+         ParLinearForm body_force_lf(&displacement_fes);
+         body_force_coef = new VectorConstantCoefficient(g);
+         auto integ = new VectorDomainLFIntegrator(*body_force_coef);
+         integ->SetIntRule(&displacement_ir);
+         body_force_lf.AddDomainIntegrator(integ);
+         body_force_lf.Assemble();
+         body_force_lf.ParallelAssemble(body_force);
+      }
+   }
+
+   void Mult(const Vector &displacement, Vector &r) const override
+   {
+      //momentum->SetParameters({mesh_nodes});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.SetSubVector(displacement_ess_tdof, 0.0);
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      jacobian_operator = std::make_shared<ElasticityJacobianOperator>(this, x);
+      return *jacobian_operator;
+
+      // fd_jacobian = std::make_shared<FDJacobian>(*this, x);
+      // return *fd_jacobian;
+   }
+
+   real_t density;
+   std::shared_ptr<DifferentiableOperator> momentum;
+   mutable std::shared_ptr<HypreParMatrix> A;
+   VectorConstantCoefficient *body_force_coef = nullptr;
+   Vector body_force;
+
+   ParGridFunction *mesh_nodes;
+
+   const Array<int> displacement_ess_tdof;
+
+   ParFiniteElementSpace &displacement_fes;
+   IntegrationRule displacement_ir;
+
+   mutable std::shared_ptr<ElasticityJacobianOperator> jacobian_operator;
+   mutable std::shared_ptr<FDJacobian> fd_jacobian;
+
+   ParFiniteElementSpace &modulus_fes;
+   ParGridFunction *modulus;
+};
+
+class StrainEnergyQoi : public Operator
+{
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+   static constexpr int ElasticModulus = 2;
+
+public:
+   StrainEnergyQoi(ParFiniteElementSpace &displacement_fes,
+                   const IntegrationRule &displacement_ir,
+                   ParGridFunction& elastic_modulus) :
+      Operator(displacement_fes.GetTrueVSize()),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      modulus_fes(*elastic_modulus.ParFESpace()),
+      modulus(&elastic_modulus)
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+      //ParametricSpace scalar_space(displacement_fes.GetMesh()->Dimension(), 0, 0, 0, 1);
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Displacement, &displacement_fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Coordinates, &mesh_fes},
+            FieldDescriptor{ElasticModulus, &modulus_fes}
+         };
+
+         energy_functional =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+         energy_functional->DisableTensorProductStructure();
+
+         mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, Weight{}, Value<ElasticModulus>{}};
+         // How do I indicate a real number output?
+         mfem::tuple outputs{One<ElasticModulus>{}};
+
+         auto qfunction = ParameterizedNeoHookeanEnergy<DIMENSION>{.nu = 0.25};
+         auto derivatives = std::integer_sequence<size_t, Displacement> {};
+         Array<int> solid_domain_attr(mesh->attributes.Max());
+         solid_domain_attr[0] = 1;
+         energy_functional->AddDomainIntegrator(
+            qfunction, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+         
+         energy_functional->SetParameters({mesh_nodes, &elastic_modulus});
+      }
+   }
+
+   void Mult(const Vector &displacement, Vector &strain_energy) const override
+   {
+      energy_functional->Mult(displacement, strain_energy);
+   }
+
+   std::shared_ptr<DifferentiableOperator> energy_functional;
+   ParGridFunction *mesh_nodes;
+   ParFiniteElementSpace &displacement_fes;
+   IntegrationRule displacement_ir;
+   ParFiniteElementSpace &modulus_fes;
+   ParGridFunction *modulus;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "/Users/andrej1/dump/fsi.msh";
+   // const char* mesh_file = "../data/ref-square.mesh";
+   int polynomial_order = 2;
+   int ir_order = 2;
+   int refinements = 0;
+   int nonlinear_solver_type = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&nonlinear_solver_type, "-nls", "--nonlinear-solver", "");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   mfem::Mesh mesh_serial = Mesh::MakeCartesian2D(20, 2, Element::QUADRILATERAL, false, 1.0,
+                                       0.1);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   mesh_serial.EnsureNodes();
+   auto mesh_beam = ParMesh(MPI_COMM_WORLD, mesh_serial);
+
+   out << "#el: " << mesh_beam.GetNE() << "\n";
+
+   H1_FECollection displacement_fec(polynomial_order, dim);
+   ParFiniteElementSpace displacement_fes(&mesh_beam, &displacement_fec, dim);
+
+   const int parameter_polynomial_order = 0;
+   L2_FECollection modulus_fec(0, dim);
+   ParFiniteElementSpace modulus_fes(&mesh_beam, &modulus_fec, dim);
+
+   HYPRE_BigInt global_size = displacement_fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "Number of unknowns: " << global_size << "\n";
+   }
+
+   const IntegrationRule &displacement_ir =
+      IntRules.Get(displacement_fes.GetFE(0)->GetGeomType(),
+                   2 * ir_order + displacement_fes.GetFE(0)->GetOrder());
+
+   Array<int> bdr_attr_is_ess(mesh_beam.bdr_attributes.Max());
+   Array<int> displacement_ess_tdof;
+   Array<int> bc_tdof;
+
+   // fixed left end
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[3] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+ 
+   ParGridFunction u(&displacement_fes);
+   u = 0.0;
+
+   ParGridFunction E(&modulus_fes);
+   E = 1.0;
+
+   ElasticityOperator elasticity(displacement_fes, displacement_ess_tdof,
+                                 displacement_ir, E);
+
+   ElasticityOperator::ElasticityJacobianPreconditioner prec;
+
+   CGSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(0.0);
+   solver.SetRelTol(1e-10);
+   // solver.SetKDim(500);
+   solver.SetMaxIter(500);
+   solver.SetPrintLevel(2);
+   solver.SetPreconditioner(prec);
+
+   std::shared_ptr<NewtonSolver> nonlinear_solver;
+   nonlinear_solver = std::make_shared<NewtonSolver>(MPI_COMM_WORLD);
+   nonlinear_solver->SetOperator(elasticity);
+   nonlinear_solver->SetRelTol(1e-9);
+   nonlinear_solver->SetMaxIter(50);
+   nonlinear_solver->SetSolver(solver);
+   nonlinear_solver->SetPrintLevel(1);
+
+   Vector zero, x(displacement_fes.GetTrueVSize());
+
+   //real_t ubc = applied_displacement(time);
+//    real_t ubc = 0.01;
+//    u.SetSubVector(bc_tdof, ubc);
+   u.GetTrueDofs(x);
+   nonlinear_solver->Mult(zero, x);
+   u.SetFromTrueDofs(x);
+   out << "Solve complete\n" << std::endl;
+
+   ParaViewDataCollection dc("dfem_elasticity_vjp", &mesh_beam);
+   dc.SetHighOrderOutput(true);
+   dc.RegisterField("displacement", &u);
+   dc.Save();
+
+   // Compute quantity of interest
+   out << "Computing QoI from solution" << std::endl;
+   StrainEnergyQoi qoi(displacement_fes, displacement_ir, E);
+   Vector energy;
+   qoi.Mult(u, energy);
+   out << "Energy = " << energy(0) << std::endl;
+
+   // TODO:
+   // 1. G = -∂(qoi)/∂u (ie adjoint load)
+   // 2. solve K^T lambda = G
+   //    In this case, K = K^T, so we can ignore the transpose matrix (or its action) action for now.
+   //    We could get it from the same trick we do in step 3.
+   // 3. d(qoi)/dE = ∂(qoi)/∂E + lambda * ∂r/∂E
+   //    To compute lambda * ∂r/∂E, set up another differentiable integrator J (u, E; v) -> reals
+   //    J(u, E; v) = v * r  (r is the residual. J will have its own q-function that looks like the virtual work)
+   //    then lambda * ∂r/∂E = ∂J(u, E; lambda)/∂E
+   // 4. Check d(qoi)/dE * dE with finite differences for some random dE vector
+
+   // To complete this, we need Julian to implement scalar-valued differentiable operators
+
+   return 0;
+}

examples/dfem_heat.cpp

@@ -0,0 +1,646 @@
+#include "dfem/dfem_refactor.hpp"
+#include "fem/pbilinearform.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+constexpr int DIMENSION = 2;
+
+template <int dim = 2>
+struct TemperatureMassQFunction
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(const real_t &T,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      return mfem::tuple{T * det(J) * w};
+   }
+};
+
+template <int dim = 2>
+struct TemperatureDiffusionQFunction
+{
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, 2> &dTdxi,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      auto invJ = inv(J);
+      auto dTdx = dTdxi * invJ;
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{dTdx * JxW};
+   }
+};
+
+class HeatOperator : public TimeDependentOperator
+{
+   static constexpr int Position = 0;
+   static constexpr int Temperature = 1;
+
+   class HeatResidual : public Operator
+   {
+   public:
+      HeatResidual(
+         HeatOperator &op,
+         const real_t &gamma,
+         const Vector &T,
+         const Vector &prevT,
+         const Vector &source,
+         const Vector &prev_source) :
+         Operator(op.Height()),
+         op(op),
+         gamma(gamma),
+         prevT(prevT),
+         source(source),
+         prev_source(prev_source),
+         z(T.Size()),
+         H1tsize(op.H1fes.GetTrueVSize()) {}
+
+      void Mult(const Vector &T, Vector &R) const override
+      {
+         auto x_gf = static_cast<ParGridFunction*>(op.H1fes.GetParMesh()->GetNodes());
+
+         R = 0.0;
+
+         op.mass->SetParameters({x_gf});
+         op.mass->Mult(T, R);
+
+         // Current F(T)
+         op.diffusion->SetParameters({x_gf});
+         op.diffusion->AddMult(T, R, gamma);
+         op.mass->AddMult(source, R, -gamma);
+
+         // Previous F(T)
+         op.diffusion->SetParameters({x_gf});
+         op.diffusion->AddMult(prevT, R, gamma);
+         op.mass->AddMult(prev_source, R, -gamma);
+
+         // Previous time stepping terms
+         op.mass->SetParameters({x_gf});
+         op.mass->AddMult(prevT, R, -1.0);
+
+         R.SetSubVector(op.temperature_ess_tdof, 0.0);
+      }
+
+      Operator& GetGradient(const Vector &u) const override
+      {
+         fd_jacobian.reset(new FDJacobian(*this, u));
+         // std::ofstream fd_jac_out("fd_jac.dat");
+         // fd_jacobian->PrintMatlab(fd_jac_out);
+         // fd_jac_out.close();
+         return *fd_jacobian;
+      }
+
+      HeatOperator &op;
+      const real_t gamma;
+
+      const int H1tsize;
+
+      Vector prevT, source, prev_source;
+      mutable Vector z;
+
+      mutable std::shared_ptr<FDJacobian> fd_jacobian;
+   };
+
+public:
+   HeatOperator(
+      ParFiniteElementSpace &H1fes,
+      Array<int> &temperature_ess_attr,
+      const IntegrationRule &ir,
+      Coefficient &temperature_exact_coeff,
+      Coefficient &source_coeff) :
+      TimeDependentOperator(H1fes.GetTrueVSize()),
+      H1fes(H1fes),
+      H1tsize(H1fes.GetTrueVSize()),
+      temperature_ess_attr(temperature_ess_attr),
+      ir(ir),
+      temperature_exact_coeff(temperature_exact_coeff),
+      T_gf(&H1fes),
+      source_coeff(source_coeff),
+      source_gf(&H1fes),
+      source_tdof(H1tsize),
+      prev_source_tdof(H1tsize),
+      prevT(H1tsize)
+   {
+      auto mesh = H1fes.GetParMesh();
+      x_gf = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *x_gf->ParFESpace();
+
+      H1fes.GetEssentialTrueDofs(temperature_ess_attr, temperature_ess_tdof);
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Temperature, &H1fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Position, &mesh_fes}
+         };
+
+         mass =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         mfem::tuple inputs{Value<Temperature>{}, Gradient<Position>{}, Weight{}};
+         mfem::tuple outputs{Value<Temperature>{}};
+
+         auto mass_qf = TemperatureMassQFunction<DIMENSION> {};
+         mass->AddDomainIntegrator(mass_qf, inputs, outputs, ir);
+      }
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Temperature, &H1fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Position, &mesh_fes}
+         };
+
+         diffusion =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         mfem::tuple inputs{Gradient<Temperature>{}, Gradient<Position>{}, Weight{}};
+         mfem::tuple outputs{Gradient<Temperature>{}};
+
+         auto diffusion_qf = TemperatureDiffusionQFunction<DIMENSION> {};
+         auto derivatives = std::integer_sequence<size_t, Temperature> {};
+         diffusion->AddDomainIntegrator(
+            diffusion_qf, inputs, outputs, ir, derivatives);
+      }
+   }
+
+   void SetTime(const real_t t)
+   {
+      TimeDependentOperator::SetTime(t);
+      temperature_exact_coeff.SetTime(t);
+      source_coeff.SetTime(t);
+   }
+
+   void Step(Vector &T, real_t &t, const real_t &dt)
+   {
+      this->SetTime(t);
+      prevT = T;
+      source_gf.ProjectCoefficient(source_coeff);
+      source_gf.GetTrueDofs(prev_source_tdof);
+
+      this->SetTime(t + dt);
+      T_gf.SetFromTrueDofs(T);
+      T_gf.ProjectBdrCoefficient(temperature_exact_coeff, temperature_ess_attr);
+      T_gf.GetTrueDofs(T);
+
+      source_gf.ProjectCoefficient(source_coeff);
+      source_gf.GetTrueDofs(source_tdof);
+
+      // Implicit midpoint
+      HeatResidual residual(*this, 0.5*dt, T, prevT, source_tdof, prev_source_tdof);
+
+      GMRESSolver krylov(MPI_COMM_WORLD);
+      krylov.SetRelTol(1e-4);
+      krylov.SetMaxIter(1000);
+      krylov.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+      NewtonSolver newton(MPI_COMM_WORLD);
+      newton.SetOperator(residual);
+      newton.SetSolver(krylov);
+      newton.SetRelTol(1e-8);
+      newton.SetMaxIter(10);
+      newton.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+      Vector zero;
+      newton.Mult(zero, T);
+
+      t += dt;
+   }
+
+   std::shared_ptr<DifferentiableOperator> mass;
+   std::shared_ptr<DifferentiableOperator> diffusion;
+
+   ParGridFunction *x_gf, source_gf, T_gf;
+
+   const Array<int> temperature_ess_attr;
+   Array<int> temperature_ess_tdof;
+
+   ParFiniteElementSpace &H1fes;
+   const int H1tsize;
+
+   Vector source_tdof, prev_source_tdof, prevT;
+
+   IntegrationRule ir;
+   Coefficient &temperature_exact_coeff;
+   Coefficient &source_coeff;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "";
+   int polynomial_order_temperature = 2;
+   int refinements = 0;
+   int problem_type = 0;
+   real_t t_final = 0.0;
+   real_t dt = 1e-3;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order_temperature, "-ot", "--order-temperature", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&t_final, "-tf", "--tf", "");
+   args.AddOption(&dt, "-dt", "--dt", "");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.EnsureNodes();
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   H1_FECollection temperature_fec(polynomial_order_temperature);
+
+   ParFiniteElementSpace H1fes(&mesh, &temperature_fec);
+
+   HYPRE_BigInt global_size_temperature = H1fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "Number of temperature unknowns: " << global_size_temperature << "\n";
+   }
+
+   const IntegrationRule &integration_rule =
+      IntRules.Get(H1fes.GetFE(0)->GetGeomType(),
+                   2 * H1fes.GetFE(0)->GetOrder() + 1);
+
+   Array<int> bdr_attr_is_ess(mesh.bdr_attributes.Max());
+   bdr_attr_is_ess = 1;
+
+   Vector T(H1fes.GetTrueVSize());
+
+   ParGridFunction T_gf(&H1fes);
+
+   auto temperature_exact = [](const Vector &coords, real_t t)
+   {
+      const real_t x = coords(0);
+      const real_t y = coords(1);
+
+      return (pow(cos(y),2) + pow(sin(x),2))/exp(2.*t);
+   };
+
+   FunctionCoefficient temperature_exact_coeff(temperature_exact);
+
+   T_gf.ProjectCoefficient(temperature_exact_coeff);
+   T_gf.GetTrueDofs(T);
+
+   auto source_term = [](const Vector &coords, real_t t)
+   {
+      const real_t x = coords(0);
+      const real_t y = coords(1);
+
+      return (-2*pow(cos(x),2))/exp(2.*t) + (2*pow(cos(y),
+                                                   2))/exp(2.*t) + (2*pow(sin(x),2))/exp(2.*t) - (2.*(pow(cos(y),2) + pow(sin(x),
+                                                         2)))/exp(2.*t) - (2*pow(sin(y),2))/exp(2.*t);
+   };
+
+   FunctionCoefficient source_term_coeff(source_term);
+
+   HeatOperator heat(H1fes, bdr_attr_is_ess, integration_rule,
+                     temperature_exact_coeff, source_term_coeff);
+
+   real_t t = 0.0;
+   out << "time step: " << dt << "\n";
+   real_t t_old;
+   bool last_step = false;
+
+   T_gf.SetFromTrueDofs(T);
+
+   ParGridFunction Terr_gf(&H1fes), Tex_gf(&H1fes);
+   Terr_gf = 0.0;
+   Tex_gf.ProjectCoefficient(temperature_exact_coeff);
+
+   ParaViewDataCollection dc("dfem_heat", &mesh);
+   dc.SetHighOrderOutput(true);
+   dc.SetLevelsOfDetail(polynomial_order_temperature);
+   dc.RegisterField("temperature", &T_gf);
+   dc.RegisterField("temperature_exact", &Tex_gf);
+   dc.RegisterField("temperature_error", &Terr_gf);
+   dc.SetCycle(0);
+   dc.SetTime(0);
+   dc.Save();
+
+   for (int ti = 1; !last_step; ti++)
+   {
+      if (t + dt >= t_final)
+      {
+         dt = t_final - t;
+         last_step = true;
+      }
+      if (Mpi::Root())
+      {
+         out << "step " << std::setw(5) << ti
+             << ",\tt = " << std::setw(5) << std::setprecision(4) << t
+             << ",\tdt = " << std::setw(5) << std::setprecision(6) << dt;
+         out << std::endl;
+      }
+
+      heat.Step(T, t, dt);
+
+      T_gf.SetFromTrueDofs(T);
+      temperature_exact_coeff.SetTime(t);
+      real_t T_l2err = T_gf.ComputeL2Error(temperature_exact_coeff);
+      if (Mpi::Root())
+      {
+         out << "|T - T_exact|_L2 = " << T_l2err << std::endl;
+      }
+
+      Tex_gf.ProjectCoefficient(temperature_exact_coeff);
+      for (int i = 0; i < Terr_gf.Size(); i++)
+      {
+         Terr_gf(i) = abs(Tex_gf(i) - T_gf(i));
+      }
+
+      if (ti % 1 == 0)
+      {
+         dc.SetCycle(ti);
+         dc.SetTime(t);
+         dc.Save();
+      }
+
+      if (Mpi::Root())
+      {
+         out << "\n" << std::endl;
+      }
+   }
+
+   return 0;
+}
+
+// class HeatOperator : public Operator
+// {
+//    static constexpr int Position = 0;
+//    static constexpr int Temperature = 1;
+
+// public:
+//    HeatOperator(
+//       ParFiniteElementSpace &H1fes,
+//       Array<int> &temperature_ess_attr,
+//       const IntegrationRule &ir,
+//       Coefficient &temperature_exact_coeff,
+//       Coefficient &source_coeff) :
+//       Operator(H1fes.GetTrueVSize()),
+//       H1fes(H1fes),
+//       H1tsize(H1fes.GetTrueVSize()),
+//       temperature_ess_attr(temperature_ess_attr),
+//       ir(ir),
+//       temperature_exact_coeff(temperature_exact_coeff),
+//       T_gf(&H1fes),
+//       source_coeff(source_coeff),
+//       source_gf(&H1fes),
+//       source_tdof(H1tsize)
+//    {
+//       auto mesh = H1fes.GetParMesh();
+//       x_gf = static_cast<ParGridFunction*>(mesh->GetNodes());
+//       ParFiniteElementSpace& mesh_fes = *x_gf->ParFESpace();
+
+//       H1fes.GetEssentialTrueDofs(temperature_ess_attr, temperature_ess_tdof);
+
+//       {
+//          {
+//             auto solutions = std::vector
+//             {
+//                FieldDescriptor{Temperature, &H1fes},
+//             };
+
+//             auto parameters = std::vector
+//             {
+//                FieldDescriptor{Position, &mesh_fes}
+//             };
+
+//             diffusion =
+//                std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+//             mfem::tuple inputs{Gradient<Temperature>{}, Gradient<Position>{}, Weight{}};
+//             mfem::tuple outputs{Gradient<Temperature>{}};
+
+//             auto diffusion_qf = TemperatureDiffusionQFunction<DIMENSION> {};
+//             auto derivatives = std::integer_sequence<size_t, Temperature> {};
+//             diffusion->AddDomainIntegrator(
+//                diffusion_qf, inputs, outputs, ir, derivatives);
+//          }
+//       }
+
+//       {
+//          ParBilinearForm diffusion(&H1fes);
+//          auto integ = new DiffusionIntegrator();
+//          integ->SetIntegrationRule(ir);
+//          diffusion.AddDomainIntegrator(integ);
+//          diffusion.Assemble();
+//          diffusion.Finalize();
+//          K.reset(diffusion.ParallelAssemble());
+//       }
+
+//       {
+//          ParLinearForm source_lf(&H1fes);
+//          source_lf.AddDomainIntegrator(new DomainLFIntegrator(source_coeff));
+//          source_lf.Assemble();
+//          source_tdof = *source_lf.ParallelAssemble();
+//       }
+//    }
+
+//    void Mult(const Vector &T, Vector &R) const override
+//    {
+//       // K->Mult(T, R);
+//       diffusion->SetParameters({x_gf});
+//       diffusion->Mult(T, R);
+//       R -= source_tdof;
+//       R.SetSubVector(temperature_ess_tdof, 0.0);
+//    }
+
+//    Operator &GetGradient(const Vector &T) const override
+//    {
+//       fd_jacobian.reset(new FDJacobian(*this, T));
+//       return *fd_jacobian;
+//    }
+
+//    ParGridFunction *x_gf, source_gf, T_gf;
+
+//    std::shared_ptr<HypreParMatrix> K;
+//    mutable std::shared_ptr<FDJacobian> fd_jacobian;
+//    std::shared_ptr<DifferentiableOperator> diffusion;
+
+//    const Array<int> temperature_ess_attr;
+//    Array<int> temperature_ess_tdof;
+
+//    ParFiniteElementSpace &H1fes;
+//    const int H1tsize;
+
+//    Vector source_tdof;
+
+//    IntegrationRule ir;
+//    Coefficient &temperature_exact_coeff;
+//    Coefficient &source_coeff;
+// };
+
+// int main(int argc, char* argv[])
+// {
+//    constexpr int dim = 2;
+
+//    Mpi::Init();
+
+//    const char* device_config = "cpu";
+//    const char* mesh_file = "";
+//    int polynomial_order_temperature = 2;
+//    int refinements = 0;
+//    int problem_type = 0;
+//    real_t t_final = 0.0;
+//    real_t dt = 1e-3;
+
+//    OptionsParser args(argc, argv);
+//    args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+//    args.AddOption(&polynomial_order_temperature, "-ot", "--order-temperature", "");
+//    args.AddOption(&refinements, "-r", "--r", "");
+//    args.AddOption(&device_config, "-d", "--device",
+//                   "Device configuration string, see Device::Configure().");
+//    args.AddOption(&t_final, "-tf", "--tf", "");
+//    args.AddOption(&dt, "-dt", "--dt", "");
+//    args.ParseCheck();
+
+//    Device device(device_config);
+//    if (Mpi::Root() == 0)
+//    {
+//       device.Print();
+//    }
+
+//    out << std::setprecision(8);
+
+//    Mesh mesh_serial = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL);
+//    MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+//    for (int i = 0; i < refinements; i++)
+//    {
+//       mesh_serial.UniformRefinement();
+//    }
+//    ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+//    mesh.EnsureNodes();
+//    mesh_serial.Clear();
+
+//    out << "#el: " << mesh.GetNE() << "\n";
+
+//    H1_FECollection temperature_fec(polynomial_order_temperature);
+
+//    ParFiniteElementSpace H1fes(&mesh, &temperature_fec);
+
+//    HYPRE_BigInt global_size_temperature = H1fes.GlobalTrueVSize();
+//    if (Mpi::Root())
+//    {
+//       out << "Number of temperature unknowns: " << global_size_temperature << "\n";
+//    }
+
+//    const IntegrationRule &integration_rule =
+//       IntRules.Get(H1fes.GetFE(0)->GetGeomType(),
+//                    2 * H1fes.GetFE(0)->GetOrder() + 1);
+
+//    Array<int> bdr_attr_is_ess(mesh.bdr_attributes.Max());
+//    bdr_attr_is_ess = 1;
+
+//    Vector T(H1fes.GetTrueVSize());
+
+//    ParGridFunction T_gf(&H1fes);
+
+//    auto temperature_exact = [](const Vector &coords, real_t t)
+//    {
+//       const real_t x = coords(0);
+//       const real_t y = coords(1);
+
+//       return pow(cos(y),2) + pow(sin(x),2);
+//    };
+
+//    FunctionCoefficient temperature_exact_coeff(temperature_exact);
+
+//    T_gf.ProjectCoefficient(temperature_exact_coeff);
+//    T_gf.GetTrueDofs(T);
+
+//    auto source_term = [](const Vector &coords, real_t t)
+//    {
+//       const real_t x = coords(0);
+//       const real_t y = coords(1);
+
+//       return -2*pow(cos(x),2) + 2*pow(cos(y),2) + 2*pow(sin(x),2) - 2*pow(sin(y),2);
+//    };
+
+//    FunctionCoefficient source_term_coeff(source_term);
+
+//    HeatOperator heat(H1fes, bdr_attr_is_ess, integration_rule,
+//                      temperature_exact_coeff, source_term_coeff);
+
+
+//    CGSolver krylov(MPI_COMM_WORLD);
+//    krylov.SetRelTol(1e-4);
+//    krylov.SetMaxIter(1000);
+//    krylov.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+//    NewtonSolver newton(MPI_COMM_WORLD);
+//    newton.SetOperator(heat);
+//    newton.SetSolver(krylov);
+//    newton.SetRelTol(1e-8);
+//    newton.SetMaxIter(50);
+//    newton.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+//    Vector zero;
+//    T_gf.ProjectBdrCoefficient(temperature_exact_coeff, bdr_attr_is_ess);
+//    T_gf.GetTrueDofs(T);
+
+//    newton.Mult(zero, T);
+
+//    T_gf.SetFromTrueDofs(T);
+
+//    ParGridFunction Terr_gf(&H1fes), Tex_gf(&H1fes);
+//    Terr_gf = 0.0;
+//    Tex_gf.ProjectCoefficient(temperature_exact_coeff);
+
+//    real_t T_l2err = T_gf.ComputeL2Error(temperature_exact_coeff);
+//    if (Mpi::Root())
+//    {
+//       out << "|T - T_exact|_L2 = " << T_l2err << std::endl;
+//    }
+
+//    Tex_gf.ProjectCoefficient(temperature_exact_coeff);
+//    for (int i = 0; i < Terr_gf.Size(); i++)
+//    {
+//       Terr_gf(i) = abs(Tex_gf(i) - T_gf(i));
+//    }
+
+//    ParaViewDataCollection dc("dfem_heat", &mesh);
+//    dc.SetHighOrderOutput(true);
+//    dc.SetLevelsOfDetail(polynomial_order_temperature);
+//    dc.RegisterField("temperature", &T_gf);
+//    dc.RegisterField("temperature_exact", &Tex_gf);
+//    dc.RegisterField("temperature_error", &Terr_gf);
+//    dc.SetCycle(0);
+//    dc.SetTime(0);
+//    dc.Save();
+// }

examples/dfem_minimal_example.cpp

@@ -0,0 +1,115 @@
+#include <tuple>
+#include <type_traits>
+#include <iostream>
+#include <enzyme/enzyme>
+
+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);
+}
+
+template <typename ... Ts>
+constexpr auto decay_types(std::tuple<Ts...> const &)
+-> std::tuple<std::remove_cv_t<std::remove_reference_t<Ts>>...>;
+
+template <typename T>
+using decay_tuple = decltype(decay_types(std::declval<T>()));
+
+template <class F> struct FunctionSignature;
+
+template <typename output_t, typename... input_ts>
+struct FunctionSignature<output_t(input_ts...)>
+{
+   using return_t = output_t;
+   using parameter_ts = std::tuple<input_ts...>;
+};
+
+template <class T> struct create_function_signature;
+
+template <typename output_t, typename T, typename... input_ts>
+struct create_function_signature<output_t (T::*)(input_ts...) const>
+{
+   using type = FunctionSignature<output_t(input_ts...)>;
+};
+
+template <typename arg_ts, std::size_t... Is>
+auto create_enzyme_args(arg_ts &args,
+                        arg_ts &shadow_args,
+                        std::index_sequence<Is...>)
+{
+   ((std::cout << std::get<Is>(shadow_args) << "\n"), ...);
+   return std::tuple<enzyme::Duplicated<decltype(std::get<Is>(args))>...>
+   {
+      { std::get<Is>(args), std::get<Is>(shadow_args) }...
+   };
+}
+
+template <typename kernel_t, typename arg_ts>
+auto fwddiff_apply_enzyme(kernel_t kernel, arg_ts &&args, arg_ts &&shadow_args)
+{
+   auto arg_indices =
+      std::make_index_sequence<std::tuple_size_v<std::remove_reference_t<arg_ts>>> {};
+
+   auto enzyme_args = create_enzyme_args(args, shadow_args, arg_indices);
+
+   using kf_return_t = typename create_function_signature<
+                       decltype(&kernel_t::operator())>::type::return_t;
+
+   std::cout << "args is " << get_type_name<decltype(args)>() << "\n\n";
+   std::cout << "enzyme_args type is " << get_type_name<decltype(enzyme_args)>() <<
+             "\n\n";
+   std::cout << "return type is " << get_type_name<decltype(kf_return_t{})>() <<
+             "\n\n";
+
+   return std::apply([&](auto &&...args)
+   {
+      return enzyme::get<0>(
+                enzyme::autodiff<enzyme::Forward>
+                (+kernel, args...));
+   },
+   enzyme_args);
+}
+
+int main()
+{
+
+   auto func = [](const double &x)
+   {
+      return x*x;
+   };
+
+   using kf_param_ts = typename create_function_signature<
+                       decltype(&decltype(func)::operator())>::type::parameter_ts;
+   using kf_output_t = typename create_function_signature<
+                       decltype(&decltype(func)::operator())>::type::return_t;
+   auto kernel_args = decay_tuple<kf_param_ts> {};
+   auto kernel_shadow_args = decay_tuple<kf_param_ts> {};
+
+   std::get<0>(kernel_args) = 3;
+   std::get<0>(kernel_shadow_args) = 1;
+   const auto res = fwddiff_apply_enzyme(func, kernel_args, kernel_shadow_args);
+   std::cout << res << " == 6\n";
+
+   return 0;
+}

examples/dfem_navier_stokes.cpp

@@ -0,0 +1,429 @@
+#include "dfem/dfem_refactor.hpp"
+#include "fem/bilininteg.hpp"
+#include "fem/intrules.hpp"
+#include "fem/pbilinearform.hpp"
+#include "linalg/hypre.hpp"
+#include "linalg/operator.hpp"
+#include "linalg/tensor.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+// -\nabla \cdot (\nabla u + p * I) -> (\nabla u + p * I, \nabla v)
+template <int dim = 2>
+class MomentumQFunction
+{
+public:
+   MomentumQFunction(const double &kinematic_viscosity,
+                     const bool &formulation) :
+      kinematic_viscosity(kinematic_viscosity),
+      formulation(formulation) {}
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<double, dim> &u,
+                   const tensor<double, dim, dim> &dudxi,
+                   const double &p,
+                   const tensor<double, dim, dim> &J,
+                   const double &w) const
+   {
+      static constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto viscous_stress = -p * I + 2.0 * kinematic_viscosity * sym(dudx);
+      auto JxW = det(J) * w * transpose(invJ);
+      if (formulation == 0)
+      {
+         return mfem::tuple{(-outer(u, u) + viscous_stress) * JxW};
+      }
+      else
+      {
+         return mfem::tuple{viscous_stress * JxW};
+      }
+   }
+
+   // TODO: this might not be ok on GPU
+   const double kinematic_viscosity;
+   const bool formulation;
+};
+
+template <int dim = 2>
+class ContinuityQFunction
+{
+public:
+   ContinuityQFunction(const int &formulation) :
+      formulation(formulation) {}
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<double, dim> &u,
+                   const tensor<double, dim, dim> &dudxi,
+                   const tensor<double, dim, dim> &J,
+                   const double &w) const
+   {
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto JxW = det(J) * w;
+      auto convective = dot(dudx, u);
+      if (formulation == 1)
+      {
+         return mfem::tuple{convective * JxW};
+      }
+      else if (formulation == 2)
+      {
+         return mfem::tuple{(convective + 0.5 * tr(dudx) * u) * JxW};
+      }
+      else if (formulation == 3)
+      {
+         // ONLY VALID FOR dim == 2
+         real_t curl_u = dudx(1, 0) - dudx(0, 1);
+         // cross product u x curl(u)
+         tensor<real_t, dim> u_cross_curl_u;
+         u_cross_curl_u(0) = u(1) * curl_u;
+         u_cross_curl_u(1) = -u(0) * curl_u;
+         return mfem::tuple{-u_cross_curl_u * JxW};
+      }
+      else if (formulation == 4)
+      {
+         return mfem::tuple{(2.0 * sym(dudx) * u + 0.5 * tr(dudx) * u) * JxW};
+      }
+   }
+
+   const int formulation;
+};
+
+class NavierStokesOperator : public Operator
+{
+   static constexpr int Velocity = 0;
+   static constexpr int Pressure = 1;
+   static constexpr int Coordinates = 2;
+
+   class NavierStokesJacobianOperator : public Operator
+   {
+   public:
+      NavierStokesJacobianOperator(
+         const NavierStokesOperator *ns,
+         const Vector &x) :
+         Operator(ns->Height()),
+         ns(ns),
+         block_op(ns->block_offsets)
+      {
+         xtmp = x;
+         BlockVector xb(xtmp.ReadWrite(), ns->block_offsets);
+
+         ParGridFunction u(&ns->velocity_fes);
+         ParGridFunction p(&ns->pressure_fes);
+
+         u.SetFromTrueDofs(xb.GetBlock(0));
+         p.SetFromTrueDofs(xb.GetBlock(1));
+
+         auto mesh_nodes = static_cast<ParGridFunction*>
+                           (ns->velocity_fes.GetParMesh()->GetNodes());
+
+         momentum_du = ns->momentum->GetDerivative(Velocity, {&u, &p}, {mesh_nodes});
+         dRdp = ns->mass_conservation;
+         dRdpT = std::make_shared<TransposeOperator>(*dRdp);
+
+         if (ns->formulation == 0)
+         {
+            block_op.SetBlock(0, 0, momentum_du.get());
+         }
+         else
+         {
+            convective_du = ns->continuity->GetDerivative(Velocity, {&u, &p}, {mesh_nodes});
+            dRdu = std::make_shared<SumOperator>(momentum_du.get(), 1.0,
+                                                 convective_du.get(), 1.0, false, false);
+            block_op.SetBlock(0, 0, dRdu.get());
+         }
+
+         block_op.SetBlock(0, 1, dRdpT.get());
+         block_op.SetBlock(1, 0, dRdp.get());
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         BlockVector xb(const_cast<double*>(x.Read()), ns->block_offsets);
+         // column elimination for essential dofs
+         xtmp = x;
+
+         BlockVector xtmpb(xtmp.ReadWrite(), ns->block_offsets);
+         xtmpb.GetBlock(0).SetSubVector(ns->vel_ess_tdofs, 0.0);
+
+         block_op.Mult(xtmpb, y);
+
+         BlockVector yb(y.ReadWrite(), ns->block_offsets);
+         for (int i = 0; i < ns->vel_ess_tdofs.Size(); i++)
+         {
+            yb.GetBlock(0)[ns->vel_ess_tdofs[i]] = xb.GetBlock(0)[ns->vel_ess_tdofs[i]];
+         }
+      }
+
+      const NavierStokesOperator *ns;
+      std::shared_ptr<Operator> momentum_du;
+      std::shared_ptr<Operator> convective_du;
+
+      std::shared_ptr<Operator> dRdu;
+      std::shared_ptr<Operator> dRdp;
+      std::shared_ptr<TransposeOperator> dRdpT;
+      BlockOperator block_op;
+
+      mutable Vector xtmp;
+   };
+
+public:
+   NavierStokesOperator(ParFiniteElementSpace &velocity_fes,
+                        ParFiniteElementSpace &pressure_fes,
+                        Array<int> &offsets,
+                        Array<int> &vel_ess_tdofs,
+                        const double &kinematic_viscosity,
+                        const IntegrationRule &velocity_ir,
+                        const IntegrationRule &pressure_ir,
+                        const int &formulation) :
+      Operator(offsets.Last()),
+      block_offsets(offsets),
+      vel_ess_tdofs(vel_ess_tdofs),
+      velocity_fes(velocity_fes),
+      pressure_fes(pressure_fes),
+      mass_conservation_form(&velocity_fes, &pressure_fes),
+      formulation(formulation)
+   {
+      auto mesh = velocity_fes.GetParMesh();
+      auto mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      auto solutions = std::vector
+      {
+         FieldDescriptor{Velocity, &velocity_fes},
+         FieldDescriptor{Pressure, &pressure_fes}
+      };
+
+      auto parameters = std::vector
+      {
+         FieldDescriptor{Coordinates, &mesh_fes}
+      };
+
+      {
+         mfem::tuple input_operators{Value<Velocity>{}, Gradient<Velocity>{}, Value<Pressure>{}, Gradient<Coordinates>{}, Weight{}};
+         mfem::tuple output_operators{Gradient<Velocity>{}};
+
+         momentum =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         auto derivatives = std::integer_sequence<size_t, Velocity> {};
+
+         MomentumQFunction<2> momentum_qf(kinematic_viscosity, formulation);
+         momentum->AddDomainIntegrator(
+            momentum_qf, input_operators, output_operators, velocity_ir,
+            derivatives);
+
+         momentum->SetParameters({mesh_nodes});
+      }
+
+      auto vdfi = new VectorDivergenceIntegrator;
+      vdfi->SetIntegrationRule(pressure_ir);
+      mass_conservation_form.AddDomainIntegrator(vdfi);
+      mass_conservation_form.Assemble();
+      mass_conservation_form.Finalize();
+      mass_conservation.reset(mass_conservation_form.ParallelAssemble());
+
+      {
+         mfem::tuple input_operators{Value<Velocity>{}, Gradient<Velocity>{}, Gradient<Coordinates>{}, Weight{}};
+         mfem::tuple output_operators{Value<Velocity>{}};
+
+         continuity =
+            std::make_shared<DifferentiableOperator>(
+               std::vector{FieldDescriptor{Velocity, &velocity_fes}}, parameters, *mesh);
+
+         auto derivatives = std::integer_sequence<size_t, Velocity> {};
+
+         ContinuityQFunction<2> convective_qf(formulation);
+         continuity->AddDomainIntegrator(
+            convective_qf, input_operators, output_operators, velocity_ir,
+            derivatives);
+
+         continuity->SetParameters({mesh_nodes});
+      }
+   }
+
+   void Mult(const Vector &x, Vector &r) const override
+   {
+      Vector xu(const_cast<double *>(x.Read()) + block_offsets[0],
+                block_offsets[1] - block_offsets[0]);
+      Vector ru(r.ReadWrite() + block_offsets[0],
+                block_offsets[1] - block_offsets[0]);
+      Vector rp(r.ReadWrite() + block_offsets[1],
+                block_offsets[2] - block_offsets[1]);
+
+      momentum->Mult(x, ru);
+
+      if (formulation != 0)
+      {
+         continuity->AddMult(xu, ru);
+      }
+
+      mass_conservation->Mult(xu, rp);
+
+      ru.SetSubVector(vel_ess_tdofs, 0.0);
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      // jacobian_operator = std::make_shared<NavierStokesJacobianOperator>(this, x);
+      // return *jacobian_operator;
+
+      fd_jacobian = std::make_shared<FDJacobian>(*this, x);
+      return *fd_jacobian;
+   }
+
+   std::shared_ptr<DifferentiableOperator> momentum;
+   std::shared_ptr<DifferentiableOperator> continuity;
+
+   ParMixedBilinearForm mass_conservation_form;
+   std::shared_ptr<Operator> mass_conservation;
+   const Array<int> block_offsets;
+   const Array<int> vel_ess_tdofs;
+
+   ParFiniteElementSpace &velocity_fes;
+   ParFiniteElementSpace &pressure_fes;
+
+   mutable std::shared_ptr<NavierStokesJacobianOperator> jacobian_operator;
+   mutable std::shared_ptr<FDJacobian> fd_jacobian;
+
+   const bool formulation;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "../data/inline-quad.mesh";
+   int polynomial_order = 2;
+   int ir_order = 2;
+   int refinements = 0;
+   double kinematic_viscosity = 1.0;
+   int formulation = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&kinematic_viscosity, "-kv", "--kv", "");
+   args.AddOption(&formulation, "-f", "--f",
+                  "Formulation:"
+                  "0 - conservative form"
+                  "1 - convective form"
+                  "2 - convective skew-symmetric form");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   H1_FECollection velocity_fec(polynomial_order, dim);
+   ParFiniteElementSpace velocity_fes(&mesh, &velocity_fec, dim);
+
+   H1_FECollection pressure_fec(polynomial_order - 1, dim);
+   ParFiniteElementSpace pressure_fes(&mesh, &pressure_fec);
+
+   out << velocity_fes.GetTrueVSize() << "\n";
+   out << pressure_fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule &velocity_ir =
+      IntRules.Get(velocity_fes.GetFE(0)->GetGeomType(),
+                   ir_order * velocity_fec.GetOrder());
+
+   const IntegrationRule &pressure_ir =
+      IntRules.Get(pressure_fes.GetFE(0)->GetGeomType(),
+                   ir_order * pressure_fec.GetOrder());
+
+   Array<int> bdr_attr_is_ess(mesh.bdr_attributes.Max());
+   bdr_attr_is_ess = 1;
+   Array<int> vel_ess_tdofs;
+   velocity_fes.GetEssentialTrueDofs(bdr_attr_is_ess, vel_ess_tdofs);
+
+   ParGridFunction u(&velocity_fes);
+   ParGridFunction p(&pressure_fes);
+
+   auto u_f = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      if (y >= 1.0)
+      {
+         u(0) = 1.0;
+      }
+      else
+      {
+         u(0) = 0.0;
+      }
+      u(1) = 0.0;
+   };
+   auto u_coef = VectorFunctionCoefficient(dim, u_f);
+
+   u.ProjectCoefficient(u_coef);
+   p = 0.0;
+
+   Array<int> block_offsets(3);
+   block_offsets[0] = 0;
+   block_offsets[1] = velocity_fes.GetTrueVSize();
+   block_offsets[2] = pressure_fes.GetTrueVSize();
+   block_offsets.PartialSum();
+
+   NavierStokesOperator navierstokes(velocity_fes, pressure_fes, block_offsets,
+                                     vel_ess_tdofs, kinematic_viscosity, velocity_ir, pressure_ir,
+                                     formulation);
+
+   BlockVector x(block_offsets), y(block_offsets);
+   u.ParallelProject(x.GetBlock(0));
+
+   GMRESSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(0.0);
+   solver.SetRelTol(1e-4);
+   solver.SetKDim(100);
+   solver.SetMaxIter(500);
+   solver.SetPrintLevel(2);
+   // solver.SetPreconditioner(prec);
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.SetOperator(navierstokes);
+   newton.SetSolver(solver);
+   newton.SetRelTol(1e-8);
+   newton.SetMaxIter(50);
+   newton.SetPrintLevel(1);
+
+   Vector zero;
+   newton.Mult(zero, x);
+
+   u.SetFromTrueDofs(x.GetBlock(0));
+   p.SetFromTrueDofs(x.GetBlock(1));
+
+   ParaViewDataCollection dc("dfem_navier_stokes", &mesh);
+   dc.SetHighOrderOutput(true);
+   dc.RegisterField("velocity", &u);
+   dc.RegisterField("pressure", &p);
+   dc.Save();
+
+   return 0;
+}

examples/dfem_navier_stokes_edac.cpp

@@ -0,0 +1,566 @@
+#include "dfem/dfem_refactor.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+void vel_ldc_ic(const Vector &coords, Vector &u)
+{
+   real_t x = coords(0);
+   real_t y = coords(1);
+
+   if (y >= 1.0)
+   {
+      u(0) = 1.0;
+   }
+   else
+   {
+      u(0) = 0.0;
+   }
+   u(1) = 0.0;
+}
+
+void vel_ldc_dt(const Vector &coords, Vector &u)
+{
+   real_t x = coords(0);
+   real_t y = coords(1);
+
+   u(0) = 0.0;
+   u(1) = 0.0;
+}
+
+void vel_shear_ic(const Vector &x, Vector &u)
+{
+   real_t xi = x(0);
+   real_t yi = x(1);
+
+   real_t rho = 80.0;
+   real_t delta = 0.05;
+
+   if (yi <= 0.5)
+   {
+      u(0) = tanh(rho * (yi - 0.25));
+   }
+   else
+   {
+      u(0) = tanh(rho * (0.75 - yi));
+   }
+
+   u(1) = delta * sin(2.0 * M_PI * (xi + 0.25));
+}
+
+void vel_2d_cyl_ic(const Vector &coords, Vector &u)
+{
+   real_t x = coords(0);
+   real_t y = coords(1);
+
+   real_t H = 0.41;
+   // real_t A = 1.5;
+   // real_t U = A * sin(M_PI * t / 8.0);
+   real_t U = 0.3;
+   u(0) = 4.0 * U * y * (H - y) / (H * H);
+   u(1) = 0.0;
+}
+
+
+void vel_2d_cyl(const Vector &coords, const real_t &t, Vector &u)
+{
+   real_t x = coords(0);
+   real_t y = coords(1);
+
+   real_t H = 0.41;
+   // real_t A = 1.5;
+   // real_t U = A * sin(M_PI * t / 8.0);
+   real_t U = 0.3;
+   if (x == 0.0)
+   {
+      u(0) = 4.0 * U * y * (H - y) / (H * H);
+   }
+   else
+   {
+      u(0) = 0.0;
+   }
+   u(1) = 0.0;
+}
+
+void vel_2d_cyl_dt(const Vector &coords, const real_t &t, Vector &u)
+{
+   real_t x = coords(0);
+   real_t y = coords(1);
+
+   real_t H = 0.41;
+   // real_t A = 1.5;
+   // real_t dUdt = 1.0 / 8.0 * A * M_PI * cos(M_PI * t / 8.0);
+   real_t dUdt = 0.0;
+   if (x == 0.0)
+   {
+      u(0) = 4.0 * dUdt * y * (H - y) / (H * H);
+   }
+   else
+   {
+      u(0) = 0.0;
+   }
+   u(1) = 0.0;
+}
+
+
+// -\nabla \cdot (\nabla u + p * I) -> (\nabla u + p * I, \nabla v)
+template <int dim = 2>
+class MomentumQFunction
+{
+public:
+   MomentumQFunction(const real_t &kinematic_viscosity) :
+      kinematic_viscosity(kinematic_viscosity) {}
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim> &u,
+                   const tensor<real_t, dim, dim> &dudxi,
+                   const real_t &p,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      static constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto viscous_stress = -p * I + kinematic_viscosity * sym(dudx);
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{(outer(u, u) - viscous_stress) * JxW};
+   }
+
+   // TODO: this might not be ok on GPU
+   const real_t kinematic_viscosity;
+};
+
+template <int dim = 2>
+class ViscousStressQFunction
+{
+public:
+   ViscousStressQFunction(const real_t &kinematic_viscosity) :
+      kinematic_viscosity(kinematic_viscosity) {}
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim, dim> &dudxi,
+                   const real_t &p,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      static constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      // auto viscous_stress = -p * I + 2.0 * kinematic_viscosity * sym(dudx);
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{-(kinematic_viscosity * dudxi) * JxW};
+   }
+
+   // TODO: this might not be ok on GPU
+   const real_t kinematic_viscosity;
+};
+
+template <int dim = 2>
+class ConvectionQFunction
+{
+public:
+   ConvectionQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim> &u,
+                   const tensor<real_t, dim, dim> &dudxi,
+                   const tensor<real_t, dim> &dpdxi,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      static constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto dpdx = dpdxi * invJ;
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{(-dot(dudx, u) - dpdx) * JxW};
+   }
+};
+
+template <int dim = 2>
+class ContinuityQFunction
+{
+public:
+   ContinuityQFunction(real_t kinematic_viscosity,
+                       real_t mach_number) :
+      kinematic_viscosity(kinematic_viscosity),
+      mach_number(mach_number) {}
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(const tensor<real_t, dim> &u,
+                   const real_t &p,
+                   const tensor<real_t, dim> &dpdxi,
+                   const tensor<real_t, dim, dim> &J,
+                   const real_t &w) const
+   {
+      auto invJ = inv(J);
+      auto dpdx = dpdxi * invJ;
+      auto JxW = det(J) * w * transpose(invJ);
+      auto r = mfem::tuple{((1.0 / (mach_number*mach_number)) * u - kinematic_viscosity * dpdx) * JxW};
+      return r;
+   }
+
+   real_t kinematic_viscosity;
+   real_t mach_number;
+};
+
+class NavierStokesOperator : public TimeDependentOperator
+{
+   static constexpr int Velocity = 0;
+   static constexpr int Pressure = 1;
+   static constexpr int Coordinates = 2;
+
+public:
+   NavierStokesOperator(ParFiniteElementSpace &velocity_fes,
+                        ParFiniteElementSpace &pressure_fes,
+                        Array<int> &offsets,
+                        Array<int> &vel_ess_bdr_attr,
+                        Array<int> &vel_ess_tdofs,
+                        const real_t &kinematic_viscosity,
+                        const IntegrationRule &velocity_ir,
+                        const IntegrationRule &pressure_ir) :
+      TimeDependentOperator(offsets.Last()),
+      block_offsets(offsets),
+      vel_ess_bdr_attr(vel_ess_bdr_attr),
+      vel_ess_tdofs(vel_ess_tdofs),
+      velocity_fes(velocity_fes),
+      pressure_fes(pressure_fes),
+      Mv_form(&velocity_fes),
+      Mp_form(&pressure_fes),
+      mv_diag_inv(velocity_fes.GetTrueVSize()),
+      mp_diag_inv(pressure_fes.GetTrueVSize()),
+      rhsu(velocity_fes.GetTrueVSize()),
+      zu(velocity_fes.GetTrueVSize()),
+      rhsp(pressure_fes.GetTrueVSize()),
+      xu_gf(&velocity_fes),
+      ku_gf(&velocity_fes),
+      xp_gf(&pressure_fes)
+   {
+      auto mesh = velocity_fes.GetParMesh();
+      auto mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      auto parameters = std::vector
+      {
+         FieldDescriptor{Coordinates, &mesh_fes}
+      };
+
+      auto solutions = std::vector
+      {
+         FieldDescriptor{Velocity, &velocity_fes},
+         FieldDescriptor{Pressure, &pressure_fes}
+      };
+
+      momentum =
+         std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+      continuity =
+         std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+      {
+         mfem::tuple input_operators{Value<Velocity>{}, Gradient<Velocity>{}, Value<Pressure>{}, Gradient<Coordinates>{}, Weight{}};
+         mfem::tuple output_operators{Gradient<Velocity>{}};
+
+         MomentumQFunction<2> momentum_qf(kinematic_viscosity);
+         momentum->AddDomainIntegrator(
+            momentum_qf, input_operators, output_operators, velocity_ir);
+
+         momentum->SetParameters({mesh_nodes});
+      }
+
+      {
+         mfem::tuple input_operators{Value<Velocity>{}, Value<Pressure>{}, Gradient<Pressure>{}, Gradient<Coordinates>{}, Weight{}};
+         mfem::tuple output_operators{Gradient<Pressure>{}};
+
+         ContinuityQFunction<2> continuity_qf(kinematic_viscosity, 0.01);
+         continuity->AddDomainIntegrator(
+            continuity_qf, input_operators, output_operators, velocity_ir);
+
+         continuity->SetParameters({mesh_nodes});
+      }
+
+      BilinearFormIntegrator *integ = new VectorMassIntegrator;
+      integ->SetIntegrationRule(velocity_ir);
+      Mv_form.AddDomainIntegrator(integ);
+      Mv_form.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+      Mv_form.Assemble();
+      Mv_form.FormSystemMatrix(vel_ess_tdofs, Mv);
+
+      MvInvPC = new OperatorJacobiSmoother(Mv_form, vel_ess_tdofs);
+
+      Array<int> outlet(velocity_fes.GetParMesh()->bdr_attributes.Max());
+      outlet = 0;
+      // outlet[1] = 1;
+      pressure_fes.GetEssentialTrueDofs(outlet, pres_ess_tdofs);
+
+      integ = new MassIntegrator;
+      integ->SetIntegrationRule(pressure_ir);
+      Mp_form.AddDomainIntegrator(integ);
+      Mp_form.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+      Mp_form.Assemble();
+
+      Mp_form.FormSystemMatrix(pres_ess_tdofs, Mp);
+
+      MpInvPC = new OperatorJacobiSmoother(Mp_form, pres_ess_tdofs);
+   }
+
+   void Mult(const Vector &x, Vector &k) const override
+   {
+      Vector xu(x.GetData() + block_offsets[0],
+                block_offsets[1] - block_offsets[0]);
+      Vector xp(x.GetData() + block_offsets[1],
+                block_offsets[2] - block_offsets[1]);
+      Vector ku(k.ReadWrite() + block_offsets[0],
+                block_offsets[1] - block_offsets[0]);
+      Vector kp(k.ReadWrite() + block_offsets[1],
+                block_offsets[2] - block_offsets[1]);
+
+      // velocity_fes.GetProlongationMatrix()->Mult(xu, xu_gf);
+      // xu_gf.ProjectBdrCoefficient(*u_coef, vel_ess_bdr_attr);
+      // velocity_fes.GetProlongationMatrix()->MultTranspose(xu_gf, xu);
+
+      // Momentum solve
+      {
+         momentum->Mult(x, rhsu);
+
+         ku_gf = 0.0;
+         ku_gf.ProjectBdrCoefficient(*dudt_coef, vel_ess_bdr_attr);
+         velocity_fes.GetProlongationMatrix()->MultTranspose(ku_gf, ku);
+
+         Mv.As<ConstrainedOperator>()->EliminateRHS(ku, rhsu);
+
+         CGSolver cg(MPI_COMM_WORLD);
+         cg.SetOperator(*Mv);
+         cg.SetPreconditioner(*MvInvPC);
+         cg.SetRelTol(1e-8);
+         cg.SetAbsTol(0.0);
+         cg.SetMaxIter(300);
+         cg.SetPrintLevel(IterativeSolver::PrintLevel().None());
+
+         cg.Mult(rhsu, ku);
+      }
+
+      // Continuity solve
+      {
+         continuity->Mult(x, rhsp);
+
+         kp = 0.0;
+         Mp.As<ConstrainedOperator>()->EliminateRHS(kp, rhsp);
+
+         CGSolver cg(MPI_COMM_WORLD);
+         cg.SetOperator(*Mp);
+         cg.SetPreconditioner(*MpInvPC);
+         cg.SetRelTol(1e-8);
+         cg.SetAbsTol(0.0);
+         cg.SetMaxIter(300);
+         cg.SetPrintLevel(IterativeSolver::PrintLevel().None());
+
+         cg.Mult(rhsp, kp);
+      }
+   }
+
+   void SetDVelocityDtDirichlet(VectorCoefficient *u, VectorCoefficient *dudt)
+   {
+      u_coef = u;
+      dudt_coef = dudt;
+   }
+
+   void SetTime(real_t t) override
+   {
+      u_coef->SetTime(t);
+      dudt_coef->SetTime(t);
+   }
+
+   std::shared_ptr<DifferentiableOperator> momentum;
+   std::shared_ptr<DifferentiableOperator> continuity;
+
+   Vector mv_diag_inv, mp_diag_inv;
+   mutable Vector rhsu, rhsp, zu;
+   mutable ParGridFunction xu_gf, ku_gf, xp_gf;
+
+   ParBilinearForm Mv_form;
+   ParBilinearForm Mp_form;
+   OperatorHandle Mv;
+   OperatorHandle Mp;
+
+   mutable OperatorJacobiSmoother *MvInvPC, *MpInvPC;
+
+   const Array<int> block_offsets;
+   const Array<int> vel_ess_tdofs;
+   Array<int> pres_ess_tdofs;
+   const Array<int> vel_ess_bdr_attr;
+
+   ParFiniteElementSpace &velocity_fes;
+   ParFiniteElementSpace &pressure_fes;
+
+   VectorCoefficient *u_coef = nullptr;
+   VectorCoefficient *dudt_coef = nullptr;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "../data/fsi.msh";
+   int polynomial_order = 2;
+   int ir_order = 2;
+   int refinements = 0;
+   real_t kinematic_viscosity = 1.0e-3;
+   real_t t_final = 1.0;
+   real_t dt = 1e-4;
+   int vis_steps = 5;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&kinematic_viscosity, "-kv", "--kv", "");
+   args.AddOption(&t_final, "-tf", "--t-final",
+                  "Final time; start time is 0.");
+   args.AddOption(&dt, "-dt", "--time-step",
+                  "Time step.");
+   args.AddOption(&vis_steps, "-vs", "--visualization-steps",
+                  "Visualize every n-th timestep.");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+   mesh_serial.EnsureNodes();
+   // GridFunction *nodes = mesh_serial.GetNodes();
+   // *nodes -= -1.0;
+   // *nodes /= 2.0;
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   // mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   H1_FECollection velocity_fec(polynomial_order, dim);
+   ParFiniteElementSpace velocity_fes(&mesh, &velocity_fec, dim);
+
+   H1_FECollection pressure_fec(polynomial_order - 1, dim);
+   ParFiniteElementSpace pressure_fes(&mesh, &pressure_fec);
+
+   out << velocity_fes.GetTrueVSize() << "\n";
+   out << pressure_fes.GetTrueVSize() << "\n";
+
+   const auto &velocity_ir = IntRules.Get(velocity_fes.GetFE(0)->GetGeomType(),
+                                          2 * polynomial_order);
+
+   const auto &pressure_ir = IntRules.Get(pressure_fes.GetFE(0)->GetGeomType(),
+                                          2 * polynomial_order);
+
+   Array<int> vel_ess_bdr_attr(mesh.bdr_attributes.Max());
+   vel_ess_bdr_attr = 1;
+   // vel_ess_bdr_attr[1] = 0; // outlet
+   // vel_ess_bdr_attr[4] = 0; // beam walls
+
+   Array<int> vel_ess_tdofs;
+   velocity_fes.GetEssentialTrueDofs(vel_ess_bdr_attr, vel_ess_tdofs);
+   // velocity_fes.GetEssentialTrueDofs(Array<int> {}, vel_ess_tdofs);
+
+   ParGridFunction u(&velocity_fes);
+   ParGridFunction p(&pressure_fes);
+
+   // auto u_ic_coef = VectorFunctionCoefficient(dim, vel_2d_cyl_ic);
+   // auto u_coef = VectorFunctionCoefficient(dim, vel_2d_cyl);
+   // auto dudt_coef = VectorFunctionCoefficient(dim, vel_2d_cyl_dt);
+
+   auto u_ic_coef = VectorFunctionCoefficient(dim, vel_ldc_ic);
+   auto u_coef = VectorFunctionCoefficient(dim, vel_ldc_ic);
+   auto dudt_coef = VectorFunctionCoefficient(dim, vel_ldc_dt);
+
+   u.ProjectCoefficient(u_ic_coef);
+   u.ProjectBdrCoefficient(u_coef, vel_ess_bdr_attr);
+   p = 0.0;
+
+   Array<int> block_offsets(3);
+   block_offsets[0] = 0;
+   block_offsets[1] = velocity_fes.GetTrueVSize();
+   block_offsets[2] = pressure_fes.GetTrueVSize();
+   block_offsets.PartialSum();
+
+   NavierStokesOperator navierstokes(velocity_fes, pressure_fes, block_offsets,
+                                     vel_ess_bdr_attr, vel_ess_tdofs, kinematic_viscosity, velocity_ir, pressure_ir);
+
+   navierstokes.SetDVelocityDtDirichlet(&u_coef, &dudt_coef);
+
+   BlockVector x(block_offsets);
+   u.ParallelProject(x.GetBlock(0));
+
+   RK3SSPSolver ode_solver;
+
+   real_t t = 0.0;
+   ode_solver.Init(navierstokes);
+
+   ParGridFunction w(&pressure_fes);
+   CurlGridFunctionCoefficient curlu(&u);
+   w.ProjectCoefficient(curlu);
+
+   ParaViewDataCollection dc("dfem_navier_stokes_edac", &mesh);
+   dc.SetHighOrderOutput(true);
+   dc.SetLevelsOfDetail(polynomial_order);
+   dc.RegisterField("velocity", &u);
+   dc.RegisterField("pressure", &p);
+   dc.RegisterField("curl_u", &w);
+   dc.SetTime(t);
+   dc.SetCycle(0);
+   dc.Save();
+
+   bool done = false;
+   for (int ti = 0; !done; )
+   {
+      real_t dt_real = std::min(dt, t_final - t);
+      ode_solver.Step(x, t, dt_real);
+      ti++;
+
+      done = (t >= t_final - 1e-8*dt);
+
+      if (done || ti % vis_steps == 0)
+      {
+         if (Mpi::Root())
+         {
+            std::cout << "time step: " << ti << ", time: " << t << std::endl;
+         }
+
+         u.SetFromTrueDofs(x.GetBlock(0));
+         p.SetFromTrueDofs(x.GetBlock(1));
+         w.ProjectCoefficient(curlu);
+
+         dc.SetTime(t);
+         dc.SetCycle(ti);
+         dc.Save();
+      }
+   }
+
+   ParMixedBilinearForm div_form(&velocity_fes, &pressure_fes);
+   div_form.AddDomainIntegrator(new VectorDivergenceIntegrator);
+   div_form.Assemble();
+
+   div_form.Mult(u, p);
+
+   out << p.Norml2() << "\n";
+
+
+   return 0;
+}

examples/dfem_nonlinear_advdiff.cpp

@@ -0,0 +1,351 @@
+#include "dfem/dfem_refactor.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+template <int dim = 2>
+class AdvDiffQFunction
+{
+public:
+   AdvDiffQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator() (
+      const double &u,
+      const tensor<real_t, dim>& dudxi,
+      const tensor<real_t, dim, dim>& J,
+      const real_t& w) const
+   {
+      auto invJ = inv(J);
+
+      // Advection
+      auto b = tensor<real_t, dim> {1.0, 1.0};
+      auto advection = -b * u;
+
+      // Diffusion
+      auto K = 1.0 / (1.0 + u*u);
+      auto diffusion = K * (dudxi * invJ);
+      // auto diffusion = dudxi * invJ;
+
+      return mfem::tuple{(advection + diffusion) * transpose(invJ) * det(J) * w};
+   }
+};
+
+real_t mms_solution(const Vector& coords)
+{
+   const double x = coords(0);
+   const double y = coords(1);
+   return pow(cos(y),2) + pow(sin(x),2);
+};
+
+real_t mms_forcing(const Vector &coords)
+{
+   const double x = coords(0);
+   const double y = coords(1);
+   return 2*cos(x)*sin(x) + (8*pow(cos(x),2)*pow(sin(x),2)*(pow(cos(y),
+                                                                2) + pow(sin(x),2)))/pow(1 + pow(pow(cos(y),2) + pow(sin(x),2),2),
+                                                                      2) - (2*pow(cos(x),2))/(1 + pow(pow(cos(y),2) + pow(sin(x),2),
+                                                                            2)) + (2*pow(cos(y),2))/(1 + pow(pow(cos(y),2) + pow(sin(x),2),
+                                                                                  2)) + (2*pow(sin(x),2))/(1 + pow(pow(cos(y),2) + pow(sin(x),2),
+                                                                                        2)) - 2*cos(y)*sin(y) + (8*pow(cos(y),2)*(pow(cos(y),2) + pow(sin(x),
+                                                                                              2))*pow(sin(y),2))/pow(1 + pow(pow(cos(y),2) + pow(sin(x),2),2),
+                                                                                                    2) - (2*pow(sin(y),2))/(1 + pow(pow(cos(y),2) + pow(sin(x),2),2));
+}
+
+template <int dim = 2>
+class AdvDiffOp : public TimeDependentOperator
+{
+   static constexpr int Concentration = 0;
+   static constexpr int Coordinates = 1;
+
+   class AdvDiffGradientOp : public Operator
+   {
+   public:
+      AdvDiffGradientOp(const AdvDiffOp &a, const Vector &x, real_t h) :
+         Operator(a.Height()),
+         a(a),
+         concentration_l(a.fes.GetTrueVSize()),
+         h(h)
+      {
+         ParGridFunction g(&a.fes, concentration_l);
+         a.fes.GetProlongationMatrix()->Mult(x, g);
+         dRdu = a.adv_diff->GetDerivative(Concentration, {&g}, {a.mesh_nodes});
+      }
+
+      void Mult(const Vector &k, Vector &y) const override
+      {
+         // column elimination for essential dofs
+         k_elim = k;
+         k_elim.SetSubVector(a.ess_tdof_list, 0.0);
+
+         dRdu->Mult(k_elim, y);
+         y *= h;
+         a.M->AddMult(k_elim, y);
+
+         for (int i = 0; i < a.ess_tdof_list.Size(); i++)
+         {
+            y[a.ess_tdof_list[i]] = k[a.ess_tdof_list[i]];
+         }
+      }
+
+      const AdvDiffOp &a;
+      mutable Vector concentration_l;
+      mutable Vector k_elim;
+      real_t h;
+      std::shared_ptr<DerivativeOperator> dRdu;
+   };
+
+   class AdvDiffResidualOp : public Operator
+   {
+   public:
+      AdvDiffResidualOp(const AdvDiffOp &a, real_t dt, const Vector &x) :
+         Operator(a.Height()),
+         dt(dt),
+         a(a),
+         x(x),
+         u(x.Size()),
+         z(x.Size())
+      {
+         HypreParMatrix A;
+         u = 0.0;
+      }
+
+      void Mult(const Vector &k, Vector &R) const override
+      {
+         u = k;
+         u *= dt;
+         u += x;
+
+         a.M->Mult(k, R);
+         a.adv_diff->AddMult(u, R);
+         R -= a.mms_forcing_rhs;
+
+         R.SetSubVector(a.ess_tdof_list, 0.0);
+      }
+
+      Operator& GetGradient(const Vector &k) const override
+      {
+         u = k;
+         u *= dt;
+         u += x;
+         jacobian.reset(new AdvDiffGradientOp(a, u, dt));
+
+         return *jacobian;
+
+         // fd_jacobian.reset(new FDJacobian(*this, k));
+         // fd_jacobian->PrintMatlab(std::cout);
+         // return *fd_jacobian;
+      }
+
+      const AdvDiffOp &a;
+      double dt;
+      Vector x;
+      mutable Vector u, z;
+
+      mutable std::shared_ptr<FDJacobian> fd_jacobian;
+
+      // AD Jacobian operator dRdu
+      mutable std::shared_ptr<AdvDiffGradientOp> jacobian;
+   };
+
+public:
+   AdvDiffOp(ParFiniteElementSpace &fes, const IntegrationRule &ir,
+             const Array<int> ess_tdof_list, bool disable_tensor_product_structure = false) :
+      TimeDependentOperator(fes.GetTrueVSize()),
+      ess_tdof_list(ess_tdof_list),
+      fes(fes),
+      Mform(&fes)
+   {
+      auto mesh = fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>
+                   (mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      auto input_operators = mfem::tuple{Value<Concentration>{}, Gradient<Concentration>{}, Gradient<Coordinates>{}, Weight{}};
+      auto output_operator = mfem::tuple{Gradient<Concentration>{}};
+
+      auto solutions = std::vector{FieldDescriptor{Concentration, &fes}};
+      auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes}};
+
+      adv_diff = std::make_unique<DifferentiableOperator>(solutions, parameters,
+                                                          *mesh);
+      auto derivatives = std::integer_sequence<size_t, Concentration> {};
+      AdvDiffQFunction<2> advdiff_qf{};
+      adv_diff->DisableTensorProductStructure(disable_tensor_product_structure);
+      adv_diff->AddDomainIntegrator(advdiff_qf, input_operators, output_operator, ir,
+                                    derivatives);
+
+      adv_diff->SetParameters({mesh_nodes});
+
+      FunctionCoefficient mms_forcing_c(mms_forcing);
+      ParLinearForm Lform(&fes);
+      Lform.AddDomainIntegrator(new DomainLFIntegrator(mms_forcing_c, &ir));
+      Lform.Assemble();
+      Lform.ParallelAssemble(mms_forcing_rhs);
+
+      Mform.AddDomainIntegrator(new MassIntegrator);
+      Mform.Assemble();
+      Mform.FormSystemMatrix(Array<int> {}, M);
+   }
+
+   void ImplicitSolve(const double dt, const Vector &x, Vector &k) override
+   {
+      auto residual = AdvDiffResidualOp(*this, dt, x);
+
+      GMRESSolver krylov(MPI_COMM_WORLD);
+      krylov.SetRelTol(1e-6);
+      krylov.SetMaxIter(1000);
+      // krylov.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+      NewtonSolver newton(MPI_COMM_WORLD);
+      newton.SetOperator(residual);
+      newton.SetSolver(krylov);
+      newton.SetRelTol(1e-12);
+      newton.SetMaxIter(10);
+      // newton.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+      Vector zero;
+      k = x;
+      k.SetSubVector(ess_tdof_list, 0.0);
+      newton.Mult(zero, k);
+   }
+
+private:
+   std::unique_ptr<DifferentiableOperator> adv_diff;
+   Array<int> ess_tdof_list;
+   ParFiniteElementSpace &fes;
+   ParGridFunction *mesh_nodes = nullptr;
+   Vector mms_forcing_rhs;
+   ParBilinearForm Mform;
+   OperatorHandle M;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "../data/inline-quad.mesh";
+   int polynomial_order = 1;
+   int ir_order = 2;
+   int refinements = 0;
+   real_t dt = 1.0;
+   real_t t_final = 1.0;
+   int vis_steps = 5;
+   bool disable_tensor_product_structure = false;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&t_final, "-tf", "--t-final", "Final time; start time is 0.");
+   args.AddOption(&dt, "-dt", "--time-step", "Time step.");
+   args.AddOption(&vis_steps, "-vs", "--visualization-steps",
+                  "Visualize every n-th timestep.");
+   args.AddOption(&disable_tensor_product_structure, "-disable-tp", "--disable-tp",
+                  "-enable-tp", "--enable-tp", "");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   out << "#dofs " << h1fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   printf("#ndof per el = %d\n", h1fes.GetFE(0)->GetDof());
+   printf("#nqp = %d\n", ir.GetNPoints());
+   printf("#q1d = %d\n", (int)floor(pow(ir.GetNPoints(), 1.0/dim) + 0.5));
+
+   ParGridFunction concentration_exact(&h1fes);
+
+   FunctionCoefficient mms_solution_c(mms_solution);
+
+   Array<int> ess_tdof_list;
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+   AdvDiffOp advdiff(h1fes, ir, ess_tdof_list, disable_tensor_product_structure);
+
+   ODESolver *ode_solver = new SDIRK23Solver;
+   ode_solver->Init(advdiff);
+
+   Vector zero, x(h1fes.GetTrueVSize());
+   concentration_exact = 0.0;
+   concentration_exact.ProjectBdrCoefficient(mms_solution_c, ess_bdr);
+   concentration_exact.GetTrueDofs(x);
+
+   // print_vector(x);
+
+   real_t t = 0.0;
+
+   bool done = false;
+   for (int ti = 0; !done; )
+   {
+      real_t dt_real = std::min(dt, t_final - t);
+      ode_solver->Step(x, t, dt_real);
+      // print_vector(x);
+      ti++;
+
+      done = (t >= t_final - 1e-8*dt);
+
+      if (done || ti % vis_steps == 0)
+      {
+         if (Mpi::Root())
+         {
+            std::cout << "time step: " << ti << ", time: " << t << std::endl;
+         }
+      }
+   }
+
+   concentration_exact.ProjectCoefficient(mms_solution_c);
+   ParGridFunction sol(&h1fes), exact_sol(&h1fes), err(&h1fes);
+   exact_sol.ProjectCoefficient(mms_solution_c);
+   sol.SetFromTrueDofs(x);
+   real_t l2err = sol.ComputeL2Error(mms_solution_c);
+   out << "l2err = " << l2err << "\n";
+
+   for (int i = 0; i < err.Size(); i++)
+   {
+      err[i] = abs(sol[i] - concentration_exact[i]);
+   }
+
+   ParaViewDataCollection dc("dfem_nonlinear_advdiff", &mesh);
+   dc.RegisterField("concentration", &sol);
+   dc.RegisterField("exact", &exact_sol);
+   dc.RegisterField("err", &err);
+   dc.Save();
+
+   return 0;
+}

examples/dfem_nonlinear_elasticity.cpp

@@ -0,0 +1,388 @@
+#include "dfem/dfem_refactor.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+constexpr int DIMENSION = 2;
+
+template <int dim = DIMENSION>
+struct MomentumRefStateQFunction
+{
+   MomentumRefStateQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudxi,
+      const tensor<real_t, dim, dim> &J,
+      const double &w) const
+   {
+      constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      constexpr real_t nu = 0.4;
+      constexpr real_t mu = 0.5 * 1e6;
+      constexpr real_t lambda = 2.0 * mu * nu / (1.0 - 2.0 * nu);
+
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto F = I + dudx;
+      // St. Venant-Kirchhoff model
+      auto C = transpose(F) * F;
+      auto E = 0.5 * (C - I);
+      auto PK2 = lambda * tr(E) * I + 2.0 * mu * E;
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{F * PK2 * JxW};
+      // auto invJ = inv(J);
+      // auto eps = sym(dudxi * invJ);
+      // return mfem::tuple{transpose(lambda * tr(eps) * I + 2.0 * mu * eps) * det(J) * w * transpose(invJ)};
+   }
+};
+
+class ElasticityOperator : public Operator
+{
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+
+public:
+   class ElasticityJacobianPreconditioner : public Solver
+   {
+   public:
+      ElasticityJacobianPreconditioner() : Solver() {}
+
+      void SetOperator(const Operator &op) override
+      {
+         this->height = op.Height();
+         this->width = op.Width();
+
+         auto elasticity_jacobian = dynamic_cast<const ElasticityJacobianOperator*>(&op);
+         MFEM_VERIFY(elasticity_jacobian != nullptr, "invalid operator");
+
+         A = std::make_shared<HypreParMatrix>();
+         elasticity_jacobian->momentum_du->Assemble(*A);
+         auto Ae = A->EliminateRowsCols(
+                      elasticity_jacobian->elasticity->displacement_ess_tdof);
+         delete Ae;
+
+         slu = std::make_shared<SuperLUSolver>(MPI_COMM_WORLD);
+         slu->SetPrintStatistics(false);
+         A_SLU = std::make_shared<SuperLURowLocMatrix>(*A);
+         slu->SetOperator(*A_SLU);
+
+         // amg = std::make_shared<HypreBoomerAMG>();
+         // amg->SetOperator(*A);
+         // amg->SetPrintLevel(0);
+         // amg->SetSystemsOptions(
+         //    elasticity_jacobian->elasticity->mesh_nodes->ParFESpace()->GetMesh()->Dimension(),
+         //    true);
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         slu->Mult(x, y);
+      }
+
+      std::shared_ptr<HypreParMatrix> A;
+      std::shared_ptr<SuperLURowLocMatrix> A_SLU;
+
+      std::shared_ptr<SuperLUSolver> slu;
+      std::shared_ptr<HypreBoomerAMG> amg;
+   };
+
+   class ElasticityJacobianOperator : public Operator
+   {
+   public:
+      ElasticityJacobianOperator(const ElasticityOperator *elasticity,
+                                 const Vector &x) :
+         Operator(elasticity->Height()),
+         elasticity(elasticity),
+         z(elasticity->Height())
+      {
+         ParGridFunction u(&elasticity->displacement_fes);
+         u.SetFromTrueDofs(x);
+
+         auto mesh_nodes = static_cast<ParGridFunction*>
+                           (elasticity->displacement_fes.GetParMesh()->GetNodes());
+         momentum_du = elasticity->momentum->GetDerivative(Displacement, {&u}, {mesh_nodes});
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         z = x;
+         z.SetSubVector(elasticity->displacement_ess_tdof, 0.0);
+
+         momentum_du->Mult(z, y);
+
+         for (int i = 0; i < elasticity->displacement_ess_tdof.Size(); i++)
+         {
+            y[elasticity->displacement_ess_tdof[i]] =
+               x[elasticity->displacement_ess_tdof[i]];
+         }
+      }
+
+      const ElasticityOperator *elasticity;
+      std::shared_ptr<DerivativeOperator> momentum_du;
+      mutable Vector z;
+   };
+
+   ElasticityOperator(ParFiniteElementSpace &displacement_fes,
+                      Array<int> &vel_ess_tdofs,
+                      const IntegrationRule &displacement_ir) :
+      Operator(displacement_fes.GetTrueVSize()),
+      density(1.0e3),
+      displacement_ess_tdof(vel_ess_tdofs),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      body_force(displacement_fes.GetTrueVSize())
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Displacement, &displacement_fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Coordinates, &mesh_fes}
+         };
+
+         momentum =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+         // momentum->DisableTensorProductStructure();
+
+         mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, Weight{}};
+         mfem::tuple outputs{Gradient<Displacement>{}};
+
+         auto momentum_qf = MomentumRefStateQFunction<DIMENSION> {};
+         auto derivatives = std::integer_sequence<size_t, Displacement> {};
+         Array<int> solid_domain_attr(mesh->attributes.Max());
+         solid_domain_attr[0] = 1;
+         momentum->AddDomainIntegrator(
+            momentum_qf, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+      }
+
+      {
+         Vector g(DIMENSION);
+         g = 0.0;
+         g(1) = 2.0 * density;
+
+         ParLinearForm body_force_lf(&displacement_fes);
+         body_force_coef = new VectorConstantCoefficient(g);
+         auto integ = new VectorDomainLFIntegrator(*body_force_coef);
+         integ->SetIntRule(&displacement_ir);
+         body_force_lf.AddDomainIntegrator(integ);
+         body_force_lf.Assemble();
+         body_force_lf.ParallelAssemble(body_force);
+      }
+   }
+
+   void Mult(const Vector &displacement, Vector &r) const override
+   {
+      momentum->SetParameters({mesh_nodes});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.SetSubVector(displacement_ess_tdof, 0.0);
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      jacobian_operator = std::make_shared<ElasticityJacobianOperator>(this, x);
+      return *jacobian_operator;
+
+      // fd_jacobian = std::make_shared<FDJacobian>(*this, x);
+      // return *fd_jacobian;
+   }
+
+   real_t density;
+   std::shared_ptr<DifferentiableOperator> momentum;
+   mutable std::shared_ptr<HypreParMatrix> A;
+   VectorConstantCoefficient *body_force_coef = nullptr;
+   Vector body_force;
+
+   ParGridFunction *mesh_nodes;
+
+   const Array<int> displacement_ess_tdof;
+
+   ParFiniteElementSpace &displacement_fes;
+   IntegrationRule displacement_ir;
+
+   mutable std::shared_ptr<ElasticityJacobianOperator> jacobian_operator;
+   mutable std::shared_ptr<FDJacobian> fd_jacobian;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "/Users/andrej1/dump/fsi.msh";
+   // const char* mesh_file = "../data/ref-square.mesh";
+   int polynomial_order = 2;
+   int ir_order = 2;
+   int refinements = 0;
+   int nonlinear_solver_type = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&nonlinear_solver_type, "-nls", "--nonlinear-solver", "");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.EnsureNodes();
+   mesh_serial.Clear();
+
+   // Array<int> beam_attributes(1);
+   // beam_attributes[0] = 2;
+   // auto mesh_beam = ParSubMesh::CreateFromDomain(mesh, beam_attributes);
+
+
+   mesh_serial = Mesh::MakeCartesian2D(8, 8, Element::QUADRILATERAL, false, 0.35,
+                                       0.02);
+   mesh_serial.EnsureNodes();
+   auto mesh_beam = ParMesh(MPI_COMM_WORLD, mesh_serial);
+
+   out << "#el: " << mesh_beam.GetNE() << "\n";
+
+   H1_FECollection displacement_fec(polynomial_order, dim);
+   ParFiniteElementSpace displacement_fes(&mesh_beam, &displacement_fec, dim);
+
+   HYPRE_BigInt global_size = displacement_fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "Number of unknowns: " << global_size << "\n";
+   }
+
+   const IntegrationRule &displacement_ir =
+      IntRules.Get(displacement_fes.GetFE(0)->GetGeomType(),
+                   2 * ir_order + displacement_fes.GetFE(0)->GetOrder());
+
+   Array<int> bdr_attr_is_ess(mesh_beam.bdr_attributes.Max());
+   out << bdr_attr_is_ess.Size() << "\n";
+   bdr_attr_is_ess = 0;
+   // bdr_attr_is_ess[6] = 1;
+   bdr_attr_is_ess[3] = 1;
+   Array<int> displacement_ess_tdof;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, displacement_ess_tdof);
+
+   ParGridFunction u(&displacement_fes);
+   // u.Randomize(1234);
+   // u.SetSubVector(displacement_ess_tdofs, 0.0);
+   u = 0.0;
+
+   ElasticityOperator elasticity(displacement_fes, displacement_ess_tdof,
+                                 displacement_ir);
+
+   ElasticityOperator::ElasticityJacobianPreconditioner prec;
+
+   CGSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(0.0);
+   solver.SetRelTol(1e-4);
+   // solver.SetKDim(500);
+   solver.SetMaxIter(500);
+   solver.SetPrintLevel(2);
+   solver.SetPreconditioner(prec);
+
+   // NewtonSolver newton(MPI_COMM_WORLD);
+   // newton.SetOperator(elasticity);
+   // newton.SetSolver(solver);
+   // newton.SetRelTol(1e-12);
+   // newton.SetMaxIter(50);
+   // newton.SetPrintLevel(1);
+
+   std::shared_ptr<NewtonSolver> nonlinear_solver;
+   if (nonlinear_solver_type == 0)
+   {
+      nonlinear_solver = std::make_shared<NewtonSolver>(MPI_COMM_WORLD);
+   }
+   // else if (nonlinear_solver_type == 1)
+   // {
+   //    nonlinear_solver = std::make_shared<KINSolver>(MPI_COMM_WORLD, KIN_LINESEARCH);
+   // }
+   else
+   {
+      MFEM_ABORT("invalid nonlinear solver type");
+   }
+   nonlinear_solver->SetOperator(elasticity);
+   nonlinear_solver->SetRelTol(1e-6);
+   nonlinear_solver->SetMaxIter(50);
+   nonlinear_solver->SetSolver(solver);
+   nonlinear_solver->SetPrintLevel(1);
+
+   Vector zero, x(displacement_fes.GetTrueVSize());
+   u.GetTrueDofs(x);
+
+   nonlinear_solver->Mult(zero, x);
+
+   u.SetFromTrueDofs(x);
+
+   // Compute Newton residual
+   Vector r(displacement_fes.GetTrueVSize());
+   elasticity.Mult(x, r);
+   r.SetSubVector(displacement_ess_tdof, 0.0);
+   double rnorm = r.Norml2();
+   out << "||F(x) - b||_2 = " << rnorm << "\n";
+
+   // Compute CG residual
+   Vector z(displacement_fes.GetTrueVSize());
+   z = x;
+   z.SetSubVector(displacement_ess_tdof, 0.0);
+   elasticity.GetGradient(x).Mult(z, r);
+   r.Neg();
+   r += elasticity.body_force;
+   for (int i = 0; i < displacement_ess_tdof.Size(); i++)
+   {
+      r[displacement_ess_tdof[i]] = x[displacement_ess_tdof[i]];
+   }
+   double cg_rnorm = r.Norml2();
+   out << "||b - Ax||_2 = " << cg_rnorm << "\n";
+   out << "||b||_2 = " << elasticity.body_force.Norml2() << "\n";
+   out << "||b - Ax||_2 / ||b||_2 = " << cg_rnorm / elasticity.body_force.Norml2()
+       << "\n";
+
+   // DenseMatrix points(dim, 1);
+   // Vector pointA(2);
+
+   // pointA(0) = 0.6;
+   // pointA(1) = 0.2;
+
+   // Array<int> elem_ids;
+   // Array<IntegrationPoint> ips;
+   // points.SetCol(0, pointA);
+   // mesh.FindPoints(points, elem_ids, ips);
+   // Vector pA(2);
+   // u.GetVectorValue(elem_ids[0], ips[0], pA);
+
+   // out << "displacement_x = " << pA(0) << "\n";
+   // out << "displacement_y = " << pA(1) << "\n";
+
+   ParaViewDataCollection dc("dfem_elasticity", &mesh_beam);
+   dc.SetHighOrderOutput(true);
+   dc.RegisterField("deformation", &u);
+   dc.Save();
+
+   return 0;
+}

examples/dfem_objective.cpp

@@ -0,0 +1,114 @@
+#include "dfem.hpp"
+
+int main(int argc, char *argv[])
+{
+   Mpi::Init();
+
+   std::cout << std::setprecision(9);
+
+   const char *mesh_file = "../data/star.mesh";
+   int polynomial_order = 1;
+   int refinements = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.ParseCheck();
+
+   Mesh mesh_serial(mesh_file, 1, 1);
+   mesh_serial.SetCurvature(1);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   const int dim = mesh_serial.Dimension();
+
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+   mesh_serial.Clear();
+
+   constexpr int vdim = 2;
+
+   // test_partial_assembly_setup_qf(mesh, 1, polynomial_order);
+   // exit(0);
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   Array<int> ess_tdof_list;
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   std::cout << "nqpts = " << ir.GetNPoints() << std::endl;
+   std::cout << "ndofs = " << h1fes.GlobalTrueVSize() << std::endl;
+
+   ParGridFunction u(&h1fes);
+   ParGridFunction g(&h1fes);
+   ParGridFunction rho(&h1fes);
+
+   auto exact_solution = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      u(0) = x*x + y;
+      u(1) = x + 0.5*y*y;
+   };
+
+   VectorFunctionCoefficient exact_solution_coeff(dim, exact_solution);
+
+   auto objective = [](tensor<double, 2> u, double rho,
+                       tensor<double, 2, 2> J,
+                       double w)
+   {
+      return sqnorm(u) * det(J) * w;
+   };
+
+   std::tuple inputs{Value{"displacement"}, Value{"density"}, Gradient{"coordinates"}, Weight{"integration_weight"}};
+   std::tuple outputs{ One{"integral"} };
+   ElementOperator objective_eop { objective, inputs, outputs };
+
+   std::vector<Field> solution_fields{{&u, "displacement"}};
+   std::vector<Field> parameter_fields{{mesh.GetNodes(), "coordinates"}, {&rho, "density"}};
+   std::vector<Field> dependent_variables{{&u, "displacement"}};
+   DifferentiableForm dop(solution_fields, parameter_fields, dependent_variables,
+                          mesh);
+
+   dop.AddElementOperator(objective_eop, ir);
+
+   u.ProjectCoefficient(exact_solution_coeff);
+   Vector zero;
+
+   Vector y(1);
+   Vector utdof;
+   u.GetTrueDofs(utdof);
+   dop.Mult(utdof, y);
+
+   // finite difference test
+   Vector dgdu(u.Size());
+   Vector fx(y);
+   out << "g: ";
+   print_vector(fx);
+   out << "\n";
+
+   for (int i = 0; i < u.Size(); i++)
+   {
+      double h = 1e-6;
+      u(i) += h;
+      dop.Mult(u, y);
+      u(i) -= h;
+      y -= fx;
+      y /= h;
+      dgdu(i) = y(0);
+   }
+
+   out << "dgdu: ";
+   print_vector(dgdu);
+
+   // Vector dgdu = dop.GetGradientWrt({&u, "displacement"});
+
+   return 0;
+}

examples/dfem_patch_test.cpp

@@ -0,0 +1,353 @@
+#include "dfem/dfem_refactor.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+constexpr int DIMENSION = 2;
+
+template <int dim = DIMENSION>
+struct MomentumRefStateQFunction
+{
+   MomentumRefStateQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudxi,
+      const tensor<real_t, dim, dim> &J,
+      const double &w) const
+   {
+      constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      constexpr real_t nu = 0.4;
+      constexpr real_t mu = 0.5 * 1e6;
+      constexpr real_t lambda = 2.0 * mu * nu / (1.0 - 2.0 * nu);
+
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto F = I + dudx;
+      // St. Venant-Kirchhoff model
+      auto C = transpose(F) * F;
+      auto E = 0.5 * (C - I);
+      auto PK2 = lambda * tr(E) * I + 2.0 * mu * E;
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{F * PK2 * JxW};
+   }
+};
+
+class ElasticityOperator : public Operator
+{
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+
+public:
+   class ElasticityJacobianPreconditioner : public Solver
+   {
+   public:
+      ElasticityJacobianPreconditioner() : Solver() {}
+
+      void SetOperator(const Operator &op) override
+      {
+         this->height = op.Height();
+         this->width = op.Width();
+
+         auto elasticity_jacobian = dynamic_cast<const ElasticityJacobianOperator*>(&op);
+         MFEM_VERIFY(elasticity_jacobian != nullptr, "invalid operator");
+
+         A = std::make_shared<HypreParMatrix>();
+         elasticity_jacobian->momentum_du->Assemble(*A);
+         auto Ae = A->EliminateRowsCols(
+                      elasticity_jacobian->elasticity->displacement_ess_tdof);
+         delete Ae;
+
+         amg = std::make_shared<HypreBoomerAMG>();
+         amg->SetOperator(*A);
+         amg->SetPrintLevel(0);
+         amg->SetSystemsOptions(
+            elasticity_jacobian->elasticity->mesh_nodes->ParFESpace()->GetMesh()->Dimension(),
+            true);
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         amg->Mult(x, y);
+      }
+
+      std::shared_ptr<HypreParMatrix> A;
+      std::shared_ptr<HypreBoomerAMG> amg;
+   };
+
+   class ElasticityJacobianOperator : public Operator
+   {
+   public:
+      ElasticityJacobianOperator(const ElasticityOperator *elasticity,
+                                 const Vector &x) :
+         Operator(elasticity->Height()),
+         elasticity(elasticity),
+         z(elasticity->Height())
+      {
+         ParGridFunction u(&elasticity->displacement_fes);
+         u.SetFromTrueDofs(x);
+
+         auto mesh_nodes = static_cast<ParGridFunction*>
+                           (elasticity->displacement_fes.GetParMesh()->GetNodes());
+         momentum_du = elasticity->momentum->GetDerivative(Displacement, {&u}, {mesh_nodes});
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         z = x;
+         z.SetSubVector(elasticity->displacement_ess_tdof, 0.0);
+
+         momentum_du->Mult(z, y);
+
+         for (int i = 0; i < elasticity->displacement_ess_tdof.Size(); i++)
+         {
+            y[elasticity->displacement_ess_tdof[i]] =
+               x[elasticity->displacement_ess_tdof[i]];
+         }
+      }
+
+      const ElasticityOperator *elasticity;
+      std::shared_ptr<DerivativeOperator> momentum_du;
+      mutable Vector z;
+   };
+
+   ElasticityOperator(ParFiniteElementSpace &displacement_fes,
+                      Array<int> &vel_ess_tdofs,
+                      const IntegrationRule &displacement_ir) :
+      Operator(displacement_fes.GetTrueVSize()),
+      density(1.0e3),
+      displacement_ess_tdof(vel_ess_tdofs),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      body_force(displacement_fes.GetTrueVSize())
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Displacement, &displacement_fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Coordinates, &mesh_fes}
+         };
+
+         momentum =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+         // momentum->DisableTensorProductStructure();
+
+         mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, Weight{}};
+         mfem::tuple outputs{Gradient<Displacement>{}};
+
+         auto momentum_qf = MomentumRefStateQFunction<DIMENSION> {};
+         auto derivatives = std::integer_sequence<size_t, Displacement> {};
+         Array<int> solid_domain_attr(mesh->attributes.Max());
+         solid_domain_attr[0] = 1;
+         momentum->AddDomainIntegrator(
+            momentum_qf, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+      }
+
+      {
+         Vector g(DIMENSION);
+         g = 0.0;
+
+         ParLinearForm body_force_lf(&displacement_fes);
+         body_force_coef = new VectorConstantCoefficient(g);
+         auto integ = new VectorDomainLFIntegrator(*body_force_coef);
+         integ->SetIntRule(&displacement_ir);
+         body_force_lf.AddDomainIntegrator(integ);
+         body_force_lf.Assemble();
+         body_force_lf.ParallelAssemble(body_force);
+      }
+   }
+
+   void Mult(const Vector &displacement, Vector &r) const override
+   {
+      momentum->SetParameters({mesh_nodes});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.SetSubVector(displacement_ess_tdof, 0.0);
+   }
+
+   void Reaction(const Vector &displacement, Vector &r) const
+   {
+      momentum->SetParameters({mesh_nodes});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.Neg();
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      jacobian_operator = std::make_shared<ElasticityJacobianOperator>(this, x);
+      return *jacobian_operator;
+
+      // fd_jacobian = std::make_shared<FDJacobian>(*this, x);
+      // return *fd_jacobian;
+   }
+
+   real_t density;
+   std::shared_ptr<DifferentiableOperator> momentum;
+   mutable std::shared_ptr<HypreParMatrix> A;
+   VectorConstantCoefficient *body_force_coef = nullptr;
+   Vector body_force;
+
+   ParGridFunction *mesh_nodes;
+
+   const Array<int> displacement_ess_tdof;
+
+   ParFiniteElementSpace &displacement_fes;
+   IntegrationRule displacement_ir;
+
+   mutable std::shared_ptr<ElasticityJacobianOperator> jacobian_operator;
+   mutable std::shared_ptr<FDJacobian> fd_jacobian;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "../data/patch2D_quads.mesh";
+   int polynomial_order = 2;
+   int ir_order = 2;
+   int refinements = 0;
+   int nonlinear_solver_type = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&nonlinear_solver_type, "-nls", "--nonlinear-solver", "");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   // Mesh mesh_serial = Mesh::MakeCartesian2D(4, 3, Element::QUADRILATERAL,
+   //    false, 1.0, 1.0);
+   Mesh mesh_serial(mesh_file);
+   mesh_serial.EnsureNodes();
+   auto mesh_beam = ParMesh(MPI_COMM_WORLD, mesh_serial);
+
+   out << "#el: " << mesh_beam.GetNE() << "\n";
+
+   H1_FECollection displacement_fec(polynomial_order, dim);
+   ParFiniteElementSpace displacement_fes(&mesh_beam, &displacement_fec, dim);
+
+   HYPRE_BigInt global_size = displacement_fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "Number of unknowns: " << global_size << "\n";
+   }
+
+   const IntegrationRule &displacement_ir =
+      IntRules.Get(displacement_fes.GetFE(0)->GetGeomType(),
+                   2 * ir_order + displacement_fes.GetFE(0)->GetOrder());
+
+   Array<int> bdr_attr_is_ess(mesh_beam.bdr_attributes.Max());
+   Array<int> displacement_ess_tdof;
+   Array<int> bc_tdof;
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[0] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 1);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[3] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[1] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   out << "Essential tdofs" << "\n";
+   displacement_ess_tdof.Print();
+
+   // Applied displacement boundary condition
+   constexpr real_t applied_displacement = 0.2;
+   ParGridFunction u(&displacement_fes);
+   u = 0.0;
+   u.SetSubVector(bc_tdof, applied_displacement);
+
+   ElasticityOperator elasticity(displacement_fes, displacement_ess_tdof,
+                                 displacement_ir);
+
+   ElasticityOperator::ElasticityJacobianPreconditioner prec;
+
+   CGSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(0.0);
+   solver.SetRelTol(1e-4);
+   // solver.SetKDim(500);
+   solver.SetMaxIter(500);
+   solver.SetPrintLevel(2);
+   solver.SetPreconditioner(prec);
+
+   auto nonlinear_solver = std::make_shared<NewtonSolver>(MPI_COMM_WORLD);
+   nonlinear_solver->SetOperator(elasticity);
+   nonlinear_solver->SetRelTol(1e-10);
+   nonlinear_solver->SetMaxIter(50);
+   nonlinear_solver->SetSolver(solver);
+   nonlinear_solver->SetPrintLevel(1);
+
+   Vector zero, x(displacement_fes.GetTrueVSize());
+   u.GetTrueDofs(x);
+
+   nonlinear_solver->Mult(zero, x);
+
+   u.SetFromTrueDofs(x);
+
+   auto exact_solution = [](const Vector& X, Vector& u) {
+      constexpr double Lx = 1.0, Ly = 1.0;
+      u(0) = X(0)/Lx*applied_displacement;
+      constexpr real_t nu = 0.4;
+      constexpr real_t mu = 0.5 * 1e6;
+      constexpr real_t E = 2*(1 + nu)*mu;
+      real_t stretch0 = 1.0 + applied_displacement/Lx;
+      real_t strain0 = 0.5*(stretch0*stretch0 - 1.0);
+      real_t strain1 = nu/(nu - 1.0)*strain0;
+      real_t stretch1 = std::sqrt(2*strain1 + 1.0);
+      u(1) = X(1)*(stretch1 - 1.0);
+   };
+   VectorFunctionCoefficient exact_solution_coef(dim, exact_solution);
+   real_t error_norm = u.ComputeL2Error(exact_solution_coef);
+   out << "Error norm = " << error_norm << std::endl;
+   if (error_norm < 1e-10) {
+      out << "[PASS]" << std::endl;
+   } else
+   {
+      out << "[FAIL]" << std::endl;
+   }
+
+   // Compute reactions
+   // Vector r(displacement_fes.GetTrueVSize());
+   // elasticity.Reaction(x, r);
+   // ParGridFunction reaction(&displacement_fes);
+   // reaction.SetFromTrueDofs(r);
+
+   // ParaViewDataCollection dc("patch_test", &mesh_beam);
+   // dc.SetHighOrderOutput(true);
+   // dc.SetLevelsOfDetail(1);
+   // dc.RegisterField("displacement", &u);
+   // dc.RegisterField("reaction", &reaction);
+   // dc.Save();
+
+   return 0;
+}

examples/dfem_plaplacian.cpp

@@ -0,0 +1,138 @@
+#include "dfem.hpp"
+
+int main(int argc, char *argv[])
+{
+   Mpi::Init();
+
+   std::cout << std::setprecision(9);
+
+   const char *mesh_file = "../data/star.mesh";
+   int polynomial_order = 1;
+   int refinements = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.ParseCheck();
+
+   Mesh mesh_serial(mesh_file, 1, 1);
+   mesh_serial.SetCurvature(1);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   const int dim = mesh_serial.Dimension();
+
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+   mesh_serial.Clear();
+
+   constexpr int vdim = 1;
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   Array<int> ess_tdof_list;
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   std::cout << "nqpts = " << ir.GetNPoints() << std::endl;
+   std::cout << "ndofs = " << h1fes.GlobalTrueVSize() << std::endl;
+
+   ParGridFunction u(&h1fes);
+
+   auto exact_solution = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      // PRESENT
+      return pow(x,2) + 0.5*x*pow(y,2);
+   };
+
+   FunctionCoefficient exact_solution_coeff(exact_solution);
+
+   auto plaplacian = [](double u,
+                        tensor<double, 2> dudxi,
+                        tensor<double, 2, 2> J,
+                        double w)
+   {
+      using mfem::internal::tensor;
+      auto dudx = dudxi * inv(J);
+      auto JxW = transpose(inv(J)) * det(J) * w;
+      // PRESENT: Implement (1+u^2) * ∇u
+      return (1.0 + u*u) * dudx * JxW;
+   };
+
+   // PRESENT: Implement descriptors
+   std::tuple input_descriptors = {Value{"potential"}, Gradient{"potential"}, Gradient{"coordinates"}, Weight{"integration_weights"}};
+   // PRESENT: Implement descriptors
+   std::tuple output_descriptors = {Gradient{"potential"}};
+
+   ElementOperator qf {plaplacian, input_descriptors, output_descriptors};
+
+   ElementOperator forcing_qf
+   {
+      [](tensor<double, 2> coords, tensor<double, 2, 2> J, double w)
+      {
+         int p = 2;
+         double x = coords(0);
+         double y = coords(1);
+         // *INDENT-OFF*
+         double mathematica_please_help_me = 2.*pow(x,2)*pow(y,2)*(pow(x,2) + 0.5*x*pow(y,2)) + 2*pow(2*x + 0.5*pow(y,2),2)*(pow(x,2) + 0.5*x*pow(y,2)) + 2*(1 + pow(pow(x,2) + 0.5*x*pow(y,2),2)) + 1.*x*(1 + pow(pow(x,2) + 0.5*x*pow(y,2),2));
+         return mathematica_please_help_me * det(J) * w;
+        // *INDENT-ON*
+      },
+      // inputs
+      std::tuple{
+         Value{"coordinates"},
+         Gradient{"coordinates"},
+         Weight{"integration_weight"}},
+      // outputs
+      std::tuple{
+         Value{"potential"}}
+   };
+
+   std::tuple list_of_qfs{qf_1, qf_2, qf_n};
+
+   std::vector<Field> solutions{{&u, "potential"}};
+   std::vector<Field> parameters{{mesh.GetNodes(), "coordinates"}};
+   DifferentiableForm dop(solutions, parameters, mesh);
+   dop.SetEssentialTrueDofs(ess_tdof_list);
+
+   auto R = dop.GetResidual(list_of_qfs, ir);
+   auto Jacobian_aka_dRdu = dop.GetDerivative<0>(list_of_qfs, ir);
+
+   // R(u) = (\grad u, \grad v) + (f, v)
+   // dop.AddElementOperator<AD::Enzyme>(qf, ir);
+   // dop.AddElementOperator<AD::None>(forcing_qf, ir);
+
+   GMRESSolver gmres(MPI_COMM_WORLD);
+   gmres.SetRelTol(1e-12);
+   gmres.SetMaxIter(5000);
+   gmres.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.SetSolver(gmres);
+   newton.SetOperator(dop);
+   newton.SetRelTol(1e-12);
+   newton.SetMaxIter(100);
+   newton.SetPrintLevel(1);
+
+   u = 1e-6;
+   u.ProjectBdrCoefficient(exact_solution_coeff, ess_bdr);
+   Vector x;
+   u.GetTrueDofs(x);
+
+   Vector zero;
+   newton.Mult(zero, x);
+
+   u.Distribute(x);
+
+   std::cout << "|u-u_ex|_L2 = " << u.ComputeL2Error(exact_solution_coeff) << "\n";
+
+   return 0;
+}

examples/dfem_plasticity.cpp

@@ -0,0 +1,585 @@
+#include "dfem/dfem_refactor.hpp"
+
+#include <fstream>
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+constexpr int DIMENSION = 2;
+
+template <typename T, int dim>
+MFEM_HOST_DEVICE inline
+tensor<T, 3, 3> tensor_to_3D(const tensor<T, dim, dim>& A)
+{
+   tensor<T, 3, 3> A3D{};
+   for (int i = 0; i < dim; i++) {
+      for (int j = 0; j < dim; j++) {
+         A3D[i][j] = A[i][j];
+      }
+   }
+   return A3D;
+}
+
+template <typename Material, int dim = DIMENSION>
+struct InternalStateQFunction
+{
+   InternalStateQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudxi,
+      const tensor<real_t, dim, dim> &J,
+      const tensor<real_t, 10> &internal_state,
+      const double &w) const
+   {
+      auto invJ = inv(J);
+      auto dudX = dudxi * invJ;
+      auto dudX3D = tensor_to_3D(dudX);
+      //auto internal_state_new = get<1>(material(dudX3D, internal_state));
+      auto [stress, internal_state_new] = material(dudX3D, internal_state);
+      // real_t vm = sqrt(1.5)*norm(dev(stress));
+      // out << vm << " " << internal_state_new[9] << std::endl;
+      return mfem::tuple{internal_state_new};
+   }
+
+   Material material;
+};
+
+template <typename Material, int dim = DIMENSION>
+struct MomentumRefStateQFunction
+{
+   MomentumRefStateQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      const tensor<real_t, dim, dim> &dudxi,
+      const tensor<real_t, dim, dim> &J,
+      const tensor<real_t, 10> &internal_state,
+      const double &w) const
+   {
+      auto invJ = inv(J);
+      auto dudX = dudxi * invJ;
+      auto dudX3D = tensor_to_3D(dudX);
+      auto [P3D, Qnew] = material(dudX3D, internal_state);
+      auto P = mfem::internal::make_tensor<dim, dim>([&P3D](int i, int j) { return P3D[i][j]; });
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{P * JxW};
+   }
+
+   Material material;
+};
+
+
+struct J2SmallStrain {
+  static constexpr int dim = 3;         ///< spatial dimension
+  static constexpr int n_internal_states = 10;
+  static constexpr double tol = 1e-10;  ///< relative tolerance on residual mag to judge convergence of return map
+
+  real_t E;                 ///< Young's modulus
+  real_t nu;                ///< Poisson's ratio
+  real_t sigma_y;           ///< Yield strength
+  real_t Hi;                ///< Isotropic hardening modulus
+  real_t density;           ///< Mass density
+
+  /// @brief variables required to characterize the hysteresis response
+  struct InternalState {
+    tensor<double, dim, dim> plastic_strain;  ///< plastic strain
+    double accumulated_plastic_strain;        ///< uniaxial equivalent plastic strain
+  };
+
+  MFEM_HOST_DEVICE inline
+  InternalState unpack_internal_state(const tensor<real_t, n_internal_states> & packed_state) const
+  {
+      // we could use type punning here to avoid copies
+      auto plastic_strain = mfem::internal::make_tensor<dim, dim>(
+         [&packed_state](int i, int j) { return packed_state[dim*i + j]; });
+      real_t accumulated_plastic_strain = packed_state[n_internal_states - 1];
+      return {plastic_strain, accumulated_plastic_strain};
+  }
+
+  MFEM_HOST_DEVICE inline
+  tensor<real_t, n_internal_states> pack_internal_state(const tensor<real_t, dim, dim> & plastic_strain, real_t accumulated_plastic_strain) const
+  {
+      tensor<real_t, n_internal_states> packed_state{};
+      for (int i = 0, ij = 0; i < dim; i++) {
+         for (int j = 0; j < dim; j++, ij++) {
+            packed_state[ij] = plastic_strain[i][j];
+         }
+      }
+      packed_state[n_internal_states - 1] = accumulated_plastic_strain;
+      return packed_state;
+  }
+
+  MFEM_HOST_DEVICE inline
+  tuple<tensor<real_t, dim, dim>, tensor<real_t, n_internal_states>>
+  operator()(const tensor<real_t, dim, dim> & dudX, const tensor<real_t, n_internal_states> & internal_state) const
+  {
+      auto I = mfem::internal::Identity<dim>();
+      const real_t K = E / (3.0 * (1.0 - 2.0 * nu));
+      const real_t G = 0.5 * E / (1.0 + nu);
+
+      auto [plastic_strain, accumulated_plastic_strain] = unpack_internal_state(internal_state);
+
+      // (i) elastic predictor
+      auto el_strain = sym(dudX) - plastic_strain;
+      auto p = K * tr(el_strain);
+      auto s = 2.0 * G * dev(el_strain);
+      auto q = sqrt(1.5) * norm(s);
+      real_t delta_eqps = 0.0;
+
+      auto flow_strength = [this](real_t eqps) { return this->sigma_y + this->Hi*eqps; };
+
+      // (ii) admissibility
+      if (q - (sigma_y + Hi*accumulated_plastic_strain) > tol*sigma_y) {
+         // (iii) return mapping
+         real_t delta_eqps = (q - sigma_y - Hi*accumulated_plastic_strain)/(3*G + Hi);
+         auto Np = 1.5 * s / q;
+         s -= 2.0 * G * delta_eqps * Np;
+         plastic_strain += delta_eqps * Np;
+         accumulated_plastic_strain += delta_eqps;
+      }
+      auto stress = s + p * I;
+      auto internal_state_new = pack_internal_state(plastic_strain, accumulated_plastic_strain);
+      return {stress, internal_state_new};
+  }
+};
+
+class ElasticityOperator : public Operator
+{
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+   static constexpr int InternalState = 2;
+
+public:
+   class ElasticityJacobianPreconditioner : public Solver
+   {
+   public:
+      ElasticityJacobianPreconditioner() : Solver() {}
+
+      void SetOperator(const Operator &op) override
+      {
+         this->height = op.Height();
+         this->width = op.Width();
+
+         auto elasticity_jacobian = dynamic_cast<const ElasticityJacobianOperator*>(&op);
+         MFEM_VERIFY(elasticity_jacobian != nullptr, "invalid operator");
+
+         A = std::make_shared<HypreParMatrix>();
+         elasticity_jacobian->momentum_du->Assemble(*A);
+         auto Ae = A->EliminateRowsCols(
+                      elasticity_jacobian->elasticity->displacement_ess_tdof);
+         delete Ae;
+
+         amg = std::make_shared<HypreBoomerAMG>();
+         amg->SetOperator(*A);
+         amg->SetPrintLevel(0);
+         amg->SetSystemsOptions(
+            elasticity_jacobian->elasticity->mesh_nodes->ParFESpace()->GetMesh()->Dimension(),
+            true);
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         amg->Mult(x, y);
+      }
+
+      std::shared_ptr<HypreParMatrix> A;
+      std::shared_ptr<HypreBoomerAMG> amg;
+   };
+
+   class ElasticityJacobianOperator : public Operator
+   {
+   public:
+      ElasticityJacobianOperator(const ElasticityOperator *elasticity,
+                                 const Vector &x) :
+         Operator(elasticity->Height()),
+         elasticity(elasticity),
+         z(elasticity->Height())
+      {
+         ParGridFunction u(&elasticity->displacement_fes);
+         u.SetFromTrueDofs(x);
+         auto mesh_nodes = static_cast<ParGridFunction*>
+                           (elasticity->displacement_fes.GetParMesh()->GetNodes());
+         momentum_du = elasticity->momentum->GetDerivative(Displacement, {&u}, {mesh_nodes, &elasticity->internal_state});
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         z = x;
+         z.SetSubVector(elasticity->displacement_ess_tdof, 0.0);
+
+         momentum_du->Mult(z, y);
+
+         for (int i = 0; i < elasticity->displacement_ess_tdof.Size(); i++)
+         {
+            y[elasticity->displacement_ess_tdof[i]] =
+               x[elasticity->displacement_ess_tdof[i]];
+         }
+      }
+
+      const ElasticityOperator *elasticity;
+      std::shared_ptr<DerivativeOperator> momentum_du;
+      mutable Vector z;
+   };
+
+   template <typename Material>
+   ElasticityOperator(ParFiniteElementSpace &displacement_fes,
+                      Array<int> &vel_ess_tdofs,
+                      const IntegrationRule &displacement_ir,
+                      ParametricFunction &internal_state,
+                      Material material) :
+      Operator(displacement_fes.GetTrueVSize()),
+      density(1.0e3),
+      displacement_ess_tdof(vel_ess_tdofs),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      internal_state(internal_state),
+      body_force(displacement_fes.GetTrueVSize())
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Displacement, &displacement_fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Coordinates, &mesh_fes},
+            FieldDescriptor{InternalState, &internal_state.space}
+         };
+
+         momentum =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+         momentum->DisableTensorProductStructure();
+
+         mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, None<InternalState>{}, Weight{}};
+         mfem::tuple outputs{Gradient<Displacement>{}};
+
+         auto momentum_qf = MomentumRefStateQFunction<Material, DIMENSION> {.material = material};
+         auto derivatives = std::integer_sequence<size_t, Displacement> {};
+         Array<int> solid_domain_attr(mesh->attributes.Max());
+         solid_domain_attr[0] = 1;
+         momentum->AddDomainIntegrator(
+            momentum_qf, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+      }
+
+      {
+         Vector g(DIMENSION);
+         g = 0.0;
+
+         ParLinearForm body_force_lf(&displacement_fes);
+         body_force_coef = new VectorConstantCoefficient(g);
+         auto integ = new VectorDomainLFIntegrator(*body_force_coef);
+         integ->SetIntRule(&displacement_ir);
+         body_force_lf.AddDomainIntegrator(integ);
+         body_force_lf.Assemble();
+         body_force_lf.ParallelAssemble(body_force);
+      }
+   }
+
+   void Mult(const Vector &displacement, Vector &r) const override
+   {
+      momentum->SetParameters({mesh_nodes, &internal_state});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.SetSubVector(displacement_ess_tdof, 0.0);
+   }
+
+   void Reaction(const Vector &displacement, Vector &r) const
+   {
+      momentum->SetParameters({mesh_nodes, &internal_state});
+      momentum->Mult(displacement, r);
+      r -= body_force;
+      r.Neg();
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      jacobian_operator = std::make_shared<ElasticityJacobianOperator>(this, x);
+      return *jacobian_operator;
+
+      // fd_jacobian = std::make_shared<FDJacobian>(*this, x);
+      // return *fd_jacobian;
+   }
+
+   real_t density;
+   std::shared_ptr<DifferentiableOperator> momentum;
+   mutable std::shared_ptr<HypreParMatrix> A;
+   VectorConstantCoefficient *body_force_coef = nullptr;
+   Vector body_force;
+
+   ParGridFunction *mesh_nodes;
+
+   const Array<int> displacement_ess_tdof;
+
+   ParFiniteElementSpace &displacement_fes;
+   IntegrationRule displacement_ir;
+
+   ParametricFunction& internal_state;
+
+   mutable std::shared_ptr<ElasticityJacobianOperator> jacobian_operator;
+   mutable std::shared_ptr<FDJacobian> fd_jacobian;
+};
+
+
+class InternalStateUpdater : public Operator
+{
+   public:
+
+   static constexpr int Displacement = 0;
+   static constexpr int Coordinates = 1;
+   static constexpr int InternalState = 2;
+   
+   template <typename Material>
+   InternalStateUpdater(ParFiniteElementSpace &displacement_fes,
+                        const IntegrationRule &displacement_ir,
+                        ParametricFunction &internal_state,
+                        Material material) :
+      Operator(displacement_fes.GetTrueVSize()),
+      displacement_fes(displacement_fes),
+      displacement_ir(displacement_ir),
+      internal_state(internal_state)
+   {
+      auto mesh = displacement_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      auto solutions = std::vector
+      {
+         FieldDescriptor{Displacement, &displacement_fes}
+      };
+
+      auto parameters = std::vector
+      {
+         FieldDescriptor{Coordinates, &mesh_fes},
+         FieldDescriptor{InternalState, &internal_state.space}
+      };
+
+      op = std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+      op->DisableTensorProductStructure();
+
+      mfem::tuple inputs{Gradient<Displacement>{}, Gradient<Coordinates>{}, None<InternalState>{}, Weight{}};
+      mfem::tuple outputs{None<InternalState>{}};
+
+      auto qfunction = InternalStateQFunction<Material, DIMENSION> {.material = material};
+      // just a placeholder for now. We want vjps wrt both displacement and old internal state eventually
+      auto derivatives = std::integer_sequence<size_t, Displacement> {};
+      Array<int> solid_domain_attr(mesh->attributes.Max());
+      solid_domain_attr[0] = 1;
+      op->AddDomainIntegrator(
+         qfunction, inputs, outputs, displacement_ir, solid_domain_attr, derivatives);
+   }
+
+   void Mult(const Vector &displacement, Vector& internal_state_new) const override
+   {
+      op->SetParameters({mesh_nodes, &internal_state});
+      op->Mult(displacement, internal_state_new);
+   }
+
+   void VjpDisplacement(ParGridFunction &u, Vector& internal_state_old, Vector& internal_state_new_bar, Vector& displacement_bar) const
+   {
+      // u, internal_state_old, internal_state_new_bar should be const
+      out << "Sizes " << "u " << u.Size() << ", qold " << internal_state_old.Size() << ", qbar " << internal_state_new_bar.Size() << ", ubar " << displacement_bar.Size() << std::endl;
+      auto grad_op = op->GetDerivative(Displacement, {&u}, {mesh_nodes, &internal_state_old});
+      grad_op->AddMultTranspose(internal_state_new_bar, displacement_bar);
+   }
+
+   ParGridFunction *mesh_nodes;
+   ParFiniteElementSpace &displacement_fes;
+   std::shared_ptr<DifferentiableOperator> op;
+   IntegrationRule displacement_ir;
+   ParametricFunction& internal_state;
+};
+
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   int polynomial_order = 1;
+   int ir_order = 2;
+   int refinements = 0;
+   int nonlinear_solver_type = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--refinements", "");
+   args.AddOption(&ir_order, "-iro", "--integration-rule-order", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.AddOption(&nonlinear_solver_type, "-nls", "--nonlinear-solver", "");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh::MakeCartesian2D(20, 2, Element::QUADRILATERAL,
+      false, 1.0, 0.1);
+   mesh_serial.EnsureNodes();
+   auto mesh_beam = ParMesh(MPI_COMM_WORLD, mesh_serial);
+
+   out << "#el: " << mesh_beam.GetNE() << "\n";
+
+   H1_FECollection displacement_fec(polynomial_order, dim);
+   ParFiniteElementSpace displacement_fes(&mesh_beam, &displacement_fec, dim);
+
+   HYPRE_BigInt global_size = displacement_fes.GlobalTrueVSize();
+   if (Mpi::Root())
+   {
+      out << "Number of unknowns: " << global_size << "\n";
+   }
+
+   const IntegrationRule &displacement_ir =
+      IntRules.Get(displacement_fes.GetFE(0)->GetGeomType(),
+                   2 * ir_order + displacement_fes.GetFE(0)->GetOrder());
+   
+   constexpr int n_internal_state_variables = 10;
+   ParametricSpace internal_state_space(dim, n_internal_state_variables, displacement_ir.GetNPoints(),
+      n_internal_state_variables*displacement_ir.GetNPoints()*mesh_beam.GetNE());
+
+   ParametricFunction internal_state(internal_state_space);
+   internal_state = 0.0;
+   ParametricFunction internal_state_old(internal_state_space);
+   internal_state_old = 0.0;
+
+   Array<int> bdr_attr_is_ess(mesh_beam.bdr_attributes.Max());
+   Array<int> displacement_ess_tdof;
+   Array<int> bc_tdof;
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[0] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 1);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[3] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   bdr_attr_is_ess = 0;
+   bdr_attr_is_ess[1] = 1;
+   displacement_fes.GetEssentialTrueDofs(bdr_attr_is_ess, bc_tdof, 0);
+   for (auto td : bc_tdof) { displacement_ess_tdof.Append(td); };
+
+   ParGridFunction u(&displacement_fes);
+   u = 0.0;
+
+   using Material = J2SmallStrain; // StVenantKirchhoff
+   Material material{.E = 1000.0, .nu = 0.25, .sigma_y = 0.53333, .Hi = 40.0, .density = 1.0};
+   // Material material{.mu = 0.5e6, .nu = 0.4};
+
+   ElasticityOperator elasticity(displacement_fes, displacement_ess_tdof,
+                                 displacement_ir, internal_state, material);
+
+   ElasticityOperator::ElasticityJacobianPreconditioner prec;
+
+   CGSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(0.0);
+   solver.SetRelTol(1e-10);
+   // solver.SetKDim(500);
+   solver.SetMaxIter(500);
+   solver.SetPrintLevel(2);
+   solver.SetPreconditioner(prec);
+
+   std::shared_ptr<NewtonSolver> nonlinear_solver;
+   if (nonlinear_solver_type == 0)
+   {
+      nonlinear_solver = std::make_shared<NewtonSolver>(MPI_COMM_WORLD);
+   }
+   // else if (nonlinear_solver_type == 1)
+   // {
+   //    nonlinear_solver = std::make_shared<KINSolver>(MPI_COMM_WORLD, KIN_LINESEARCH);
+   // }
+   else
+   {
+      MFEM_ABORT("invalid nonlinear solver type");
+   }
+   nonlinear_solver->SetOperator(elasticity);
+   nonlinear_solver->SetRelTol(1e-9);
+   nonlinear_solver->SetMaxIter(25);
+   nonlinear_solver->SetSolver(solver);
+   nonlinear_solver->SetPrintLevel(1);
+
+   // variables for output
+   QuadratureSpace output_internal_state_space(mesh_beam, displacement_ir);
+   QuadratureFunction output_internal_state(&output_internal_state_space, internal_state.GetData(), material.n_internal_states);
+   Vector r(displacement_fes.GetTrueVSize());
+   ParGridFunction reaction(&displacement_fes);
+   Vector end_forces_x(bc_tdof.Size());
+
+   ParaViewDataCollection dc("dfem_plasticity", &mesh_beam);
+   dc.SetHighOrderOutput(true);
+   dc.SetLevelsOfDetail(1);
+   dc.RegisterField("displacement", &u);
+   dc.RegisterField("reaction", &reaction);
+   dc.RegisterQField("internal_state", &output_internal_state);
+   dc.SetCycle(0);
+   dc.Save();
+
+   InternalStateUpdater internal_state_update(displacement_fes, displacement_ir, internal_state, material);
+   //Vector q(internal_state_space.GetTotalSize());
+   
+   auto applied_displacement = [](double t) { return 1.2e-2*t; };
+
+   real_t time = 0.0;
+   std::ofstream history_file("history_output.csv");
+   history_file << applied_displacement(time) << " " << 0.0 << std::endl;
+
+   Vector zero, x(displacement_fes.GetTrueVSize());
+
+   constexpr int max_cycles = 3;
+   const real_t dt = 1.0/(max_cycles - 1);
+   for (int cycle = 1; cycle < max_cycles; cycle++) {
+      time += dt;
+      out << "-------------------------------------------" << std::endl;
+      out << "TIME STEP " << cycle << std::endl;
+      out << "t = " << time << std::endl;
+      
+      real_t ubc = applied_displacement(time);
+      u.SetSubVector(bc_tdof, ubc);
+
+      u.GetTrueDofs(x);
+      nonlinear_solver->Mult(zero, x);
+      u.SetFromTrueDofs(x);
+
+      // update internal variables
+      internal_state_old.Set(1.0, internal_state);
+      internal_state_update.Mult(u, internal_state);
+
+      // Compute reactions
+      elasticity.Reaction(x, r);
+      reaction.SetFromTrueDofs(r);
+      reaction.GetSubVector(bc_tdof, end_forces_x);
+      real_t force = -end_forces_x.Sum();
+      out << "u = " << applied_displacement(time) << ", Force = " << force << std::endl;
+      history_file << applied_displacement(time) << " " << force << std::endl;
+
+      dc.SetCycle(cycle);
+      dc.SetTime(time);
+      dc.Save();
+   }
+
+   // try to use the derivative to see if it works
+   ParametricFunction internal_state_bar(internal_state_space);
+   internal_state_bar = 1.0;
+   //ParGridFunction u_bar(displacement_fes);
+   Vector u_bar(displacement_fes.GetTrueVSize());
+   internal_state_update.VjpDisplacement(u, internal_state_old, internal_state_bar, u_bar);
+   u_bar.Print();
+
+   history_file.close();
+   return 0;
+}

examples/dfem_poisson.cpp

@@ -0,0 +1,192 @@
+#include "dfem/dfem_refactor.hpp"
+#include "linalg/hypre.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+template <typename diffusion_t, typename force_t>
+class DiffusionOperator : public Operator
+{
+   template <typename diffusion_du_t>
+   class DiffusionJacobianOperator : public Operator
+   {
+   public:
+      DiffusionJacobianOperator(const DiffusionOperator *diffusion,
+                                std::shared_ptr<diffusion_du_t> diff_du) :
+         Operator(diffusion->Height()), s(diffusion)
+      {
+         diff_du->Assemble(A);
+         A.EliminateBC(s->ess_tdofs, Operator::DiagonalPolicy::DIAG_ONE);
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         A.Mult(x, y);
+      }
+
+      const DiffusionOperator *s;
+      HypreParMatrix A;
+   };
+
+public:
+   DiffusionOperator(diffusion_t &diffusion, force_t &force,
+                     Array<int> &ess_tdofs) :
+      Operator(diffusion.Height()), diffusion(diffusion),
+      force(force), ess_tdofs(ess_tdofs), f(force.Height()) {}
+
+   void SetParameters(ParGridFunction &mesh_nodes)
+   {
+      diffusion.SetParameters({&mesh_nodes});
+      force.SetParameters({&mesh_nodes});
+
+      Vector zero;
+
+      this->mesh_nodes.SetSpace(mesh_nodes.ParFESpace());
+      this->mesh_nodes = mesh_nodes;
+   }
+
+   void Mult(const Vector &x, Vector &r) const override
+   {
+      diffusion.Mult(x, r);
+      force.Mult(x, f);
+      r -= f;
+      r.SetSubVector(ess_tdofs, 0.0);
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      ParGridFunction u(const_cast<ParFiniteElementSpace *>
+                        (*std::get_if<const ParFiniteElementSpace *>
+                         (&diffusion.solutions[0].data)));
+      u.SetFromTrueDofs(x);
+
+      auto dfdu = diffusion.template GetDerivativeWrt<0>({&u}, {&mesh_nodes});
+      dfdu->Assemble(A);
+      A.EliminateBC(ess_tdofs, DiagonalPolicy::DIAG_ONE);
+      return A;
+      // delete jacobian_operator;
+      // jacobian_operator = new
+      // DiffusionJacobianOperator<typename std::remove_pointer<decltype(dfdu.get())>::type>
+      // (this, dfdu);
+      // return *jacobian_operator;
+   }
+
+   diffusion_t &diffusion;
+   force_t &force;
+
+   const Array<int> ess_tdofs;
+   mutable Vector f;
+
+   mutable ParGridFunction mesh_nodes;
+
+   mutable Operator *jacobian_operator = nullptr;
+   mutable HypreParMatrix A;
+};
+
+int main(int argc, char *argv[])
+{
+   Mpi::Init();
+   int num_procs = Mpi::WorldSize();
+   int myid = Mpi::WorldRank();
+   Hypre::Init();
+
+   const char *mesh_file = "../data/ref-square.mesh";
+   int polynomial_order = 2;
+   int ir_order = 2;
+   int refinements = 4;
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection potential_fec(polynomial_order, dim);
+   ParFiniteElementSpace potential_fes(&mesh, &potential_fec);
+
+   const IntegrationRule &potential_ir =
+      IntRules.Get(potential_fes.GetFE(0)->GetGeomType(),
+                   ir_order * potential_fec.GetOrder());
+
+   Array<int> bdr_attr_is_ess(mesh.bdr_attributes.Max());
+   bdr_attr_is_ess = 1;
+   Array<int> ess_tdofs;
+   potential_fes.GetEssentialTrueDofs(bdr_attr_is_ess, ess_tdofs);
+
+   ParGridFunction u(&potential_fes);
+   u = 0.0;
+
+   auto diffusion_kernel = [](const internal::dual<double, double> &u,
+                              const tensor<internal::dual<double, double>, 2> &dudxi,
+                              const tensor<double, 2, 2> &J,
+                              const double &w)
+   {
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      return std::tuple{(1.0 + u * u) * dudx * det(J) * w * transpose(invJ)};
+   };
+
+   std::tuple argument_operators_0{Value{"potential"}, Gradient{"potential"}, Gradient{"coordinates"}, Weight{"integration_weights"}};
+   std::tuple output_operator_0{Gradient{"potential"}};
+   ElementOperator op_0{diffusion_kernel, argument_operators_0, output_operator_0};
+
+   auto force_kernel = [](const tensor<double, 2, 2> &J,
+                          const double &w)
+   {
+      return std::tuple{1.0 * det(J) * w};
+   };
+   std::tuple argument_operators_1{Gradient{"coordinates"}, Weight{"integration_weights"}};
+   std::tuple output_operator_1{Value{"potential"}};
+   ElementOperator op_1{force_kernel, argument_operators_1, output_operator_1};
+
+   std::array solutions{FieldDescriptor{&potential_fes, "potential"}};
+   std::array parameters{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator diffusion_op{solutions, parameters, std::tuple{op_0}, mesh, potential_ir};
+   DifferentiableOperator force_op{solutions, parameters, std::tuple{op_1}, mesh, potential_ir};
+
+   DiffusionOperator diffusion(diffusion_op, force_op, ess_tdofs);
+
+   diffusion.SetParameters({*mesh_nodes});
+
+   HypreBoomerAMG amg;
+   amg.SetPrintLevel(0);
+
+   CGSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(1e-12);
+   solver.SetRelTol(1e-12);
+   solver.SetMaxIter(500);
+   solver.SetPrintLevel(2);
+   solver.SetPreconditioner(amg);
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.SetOperator(diffusion);
+   newton.SetSolver(solver);
+   newton.SetRelTol(1e-8);
+   newton.SetMaxIter(10);
+   newton.SetPrintLevel(1);
+
+   Vector zero;
+   Vector x(potential_fes.GetTrueVSize());
+   u.ParallelProject(x);
+   newton.Mult(zero, x);
+
+   u.SetFromTrueDofs(x);
+
+   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" << mesh << u << std::flush;
+
+   return 0;
+}

examples/dfem_refactor.cpp

@@ -0,0 +1,102 @@
+#include "mfem.hpp"
+#include "dfem/dfem_refactor.hpp"
+
+using namespace mfem;
+
+auto main(int argc, char *argv[]) -> int
+{
+   Mpi::Init();
+
+   std::cout << std::setprecision(9);
+
+   const char *mesh_file = "../data/star.mesh";
+   int polynomial_order = 1;
+   int refinements = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.ParseCheck();
+
+   Mesh mesh_serial(mesh_file, 1, 1);
+   mesh_serial.SetCurvature(1);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   const int dim = mesh_serial.Dimension();
+
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   constexpr int vdim = 1;
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   Array<int> ess_tdof_list;
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   std::cout << "nqpts = " << ir.GetNPoints() << std::endl;
+   std::cout << "ndofs = " << h1fes.GlobalTrueVSize() << std::endl;
+
+   ParGridFunction u(&h1fes);
+
+   auto exact_solution = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      return 2.345 + x + y;
+   };
+
+   FunctionCoefficient exact_solution_coeff(exact_solution);
+
+   u.ProjectCoefficient(exact_solution_coeff);
+
+   auto domain_qf = [](const double &u,
+                       const tensor<double, 2, 2> &J,
+                       const double &w)
+   {
+      out << u << "\n" << J << "\n" << w << "\n\n";
+      return std::tuple{u * det(J) * w};
+   };
+
+   std::tuple input_descriptors = {Value{"potential"}, Gradient{"coordinates"}, Weight{"integration_weights"}};
+   std::tuple output_descriptors = {Value{"potential"}};
+   ElementOperator eop{domain_qf, input_descriptors, output_descriptors};
+
+   auto ops = std::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "potential"}};
+   auto parameters = std::array{FieldDescriptor{&mesh_fes, "coordinates"}};
+   DifferentiableOperator dop{solutions, parameters, ops, mesh, ir};
+
+   Vector x(h1fes.GetTrueVSize()), y(h1fes.GetTrueVSize());
+
+   u.GetTrueDofs(x);
+
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+
+   // Derivative wrt "potential", indicated by the index 0 of the set {solutions} \cup {parameters}
+   auto dFd0 = dop.GetDerivativeWrt<0>({&u}, {mesh_nodes});
+   dFd0->Mult(x, y);
+
+   Vector dFd0_vec;
+   dFd0->Assemble(dFd0_vec);
+
+   // Derivative wrt "coordinates", indicated by the index 1 of the set {solutions} \cup {parameters}
+   auto dFd1 = dop.GetDerivativeWrt<1>({&u}, {mesh_nodes});
+   dFd1->Mult(x, y);
+
+   return 0;
+}

examples/dfem_stokes.cpp

@@ -0,0 +1,375 @@
+#include "dfem/dfem_refactor.hpp"
+#include "fem/pgridfunc.hpp"
+#include "linalg/hypre.hpp"
+#include "linalg/solvers.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+template <int dim = 2>
+struct StokesMomentumQFunction
+{
+   StokesMomentumQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      // velocity gradient in reference space
+      const tensor<real_t, dim, dim> &dudxi, // internal::dual<real_t, real_t>
+      const real_t &p, // internal::dual<real_t, real_t>
+      const tensor<real_t, dim, dim> &J,
+      const double &w) const
+   {
+      constexpr real_t kinematic_viscosity = 1.0;
+      auto I = mfem::internal::IsotropicIdentity<dim>();
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto viscous_stress = -p * I + 2.0 * kinematic_viscosity * sym(dudx);
+      auto JxW = det(J) * w * transpose(invJ);
+      return mfem::tuple{-viscous_stress * JxW};
+   }
+};
+
+template <int dim = 2, int sdim = 2>
+struct StokesMassConservationQFunction
+{
+   StokesMassConservationQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator()(
+      // velocity gradient in reference space
+      const tensor<double, dim, dim> &dudxi,
+      const tensor<double, sdim, dim> &J,
+      const double &w) const
+   {
+      return mfem::tuple{tr(dudxi * inv(J)) * det(J) * w};
+   }
+};
+
+class StokesOperator : public Operator
+{
+   static constexpr int Velocity = 0;
+   static constexpr int Pressure = 1;
+   static constexpr int Coordinates = 2;
+
+   class StokesJacobianOperator : public Operator
+   {
+   public:
+      StokesJacobianOperator(const StokesOperator *ns, const Vector &x) :
+         Operator(ns->Height()),
+         ns(ns),
+         block_op(ns->block_offsets)
+      {
+         xtmp = x;
+         BlockVector xb(xtmp.ReadWrite(), ns->block_offsets);
+
+         ParGridFunction u(&ns->velocity_fes);
+         ParGridFunction p(&ns->pressure_fes);
+
+         u.SetFromTrueDofs(xb.GetBlock(0));
+         p.SetFromTrueDofs(xb.GetBlock(1));
+
+         auto mesh_nodes = static_cast<ParGridFunction*>
+                           (ns->velocity_fes.GetParMesh()->GetNodes());
+         momentum_du = ns->momentum->GetDerivative(Velocity, {&u}, {&p, mesh_nodes});
+
+         // Get a HypreParMatrix
+         //
+         // HypreParMatrix A;
+         // static_cast<DerivativeOperator *>(momentum_du.get())->Assemble(A);
+         //
+         // or directly
+         //
+         // HypreParMatrix A;
+         // ns->momentum->GetDerivative(Velocity, {&u,}, {&p, mesh_nodes})->Assemble(A);
+
+         dRdp = ns->mass_conservation;
+         dRdpT = std::make_shared<TransposeOperator>(*dRdp);
+
+         block_op.SetBlock(0, 0, momentum_du.get());
+         block_op.SetBlock(0, 1, dRdpT.get());
+         block_op.SetBlock(1, 0, dRdp.get());
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         BlockVector xb(const_cast<double*>(x.Read()), ns->block_offsets);
+         // column elimination for essential dofs
+         xtmp = x;
+
+         BlockVector xtmpb(xtmp.ReadWrite(), ns->block_offsets);
+         xtmpb.GetBlock(0).SetSubVector(ns->vel_ess_tdofs, 0.0);
+
+         block_op.Mult(xtmpb, y);
+
+         BlockVector yb(y.ReadWrite(), ns->block_offsets);
+         for (int i = 0; i < ns->vel_ess_tdofs.Size(); i++)
+         {
+            yb.GetBlock(0)[ns->vel_ess_tdofs[i]] = xb.GetBlock(0)[ns->vel_ess_tdofs[i]];
+         }
+      }
+
+      const StokesOperator *ns;
+      std::shared_ptr<Operator> momentum_du;
+      std::shared_ptr<Operator> convective_du;
+
+      std::shared_ptr<Operator> dRdu;
+      std::shared_ptr<Operator> dRdp;
+      std::shared_ptr<TransposeOperator> dRdpT;
+      BlockOperator block_op;
+
+      mutable Vector xtmp;
+   };
+
+public:
+   StokesOperator(ParFiniteElementSpace &velocity_fes,
+                  ParFiniteElementSpace &pressure_fes,
+                  Array<int> &offsets,
+                  Array<int> &vel_ess_tdofs,
+                  const IntegrationRule &velocity_ir,
+                  const IntegrationRule &pressure_ir) :
+      Operator(offsets.Last()),
+      block_offsets(offsets),
+      vel_ess_tdofs(vel_ess_tdofs),
+      velocity_fes(velocity_fes),
+      pressure_fes(pressure_fes),
+      mass_conservation_form(&velocity_fes, &pressure_fes)
+   {
+      auto mesh = velocity_fes.GetParMesh();
+      mesh_nodes = static_cast<ParGridFunction*>(mesh->GetNodes());
+      ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+      {
+         auto solutions = std::vector
+         {
+            FieldDescriptor{Velocity, &velocity_fes},
+         };
+
+         auto parameters = std::vector
+         {
+            FieldDescriptor{Pressure, &pressure_fes},
+            FieldDescriptor{Coordinates, &mesh_fes}
+         };
+
+         momentum =
+            std::make_shared<DifferentiableOperator>(solutions, parameters, *mesh);
+
+         mfem::tuple inputs{Gradient<Velocity>{}, Value<Pressure>{}, Gradient<Coordinates>{}, Weight{}};
+         mfem::tuple outputs{Gradient<Velocity>{}};
+
+         auto stokes_momemtum_qf = StokesMomentumQFunction{};
+         auto derivatives = std::integer_sequence<size_t, Velocity> {};
+         momentum->AddDomainIntegrator(stokes_momemtum_qf, inputs, outputs, velocity_ir,
+                                       derivatives);
+      }
+
+      // Standard MFEM integrator
+      auto vdfi = new VectorDivergenceIntegrator;
+      vdfi->SetIntegrationRule(pressure_ir);
+      mass_conservation_form.AddDomainIntegrator(vdfi);
+      mass_conservation_form.Assemble();
+      mass_conservation_form.Finalize();
+      mass_conservation.reset(mass_conservation_form.ParallelAssemble());
+
+      // dFEM
+      // {
+      //    auto solutions = std::vector
+      //    {
+      //       FieldDescriptor{Pressure, &pressure_fes},
+      //    };
+
+      //    auto parameters = std::vector
+      //    {
+      //       FieldDescriptor{Velocity, &velocity_fes},
+      //       FieldDescriptor{Coordinates, &mesh_fes}
+      //    };
+
+      //    mass_conservation = std::make_shared<DifferentiableOperator>(solutions,
+      //                                                                 parameters,
+      //                                                                 mesh);
+
+      //    mfem::tuple inputs{Gradient<Velocity>{}, Gradient<Coordinates>{}, Weight{}};
+      //    mfem::tuple outputs{Value<Pressure>{}};
+
+      //    auto stokes_mass_conservation_qf = StokesMassConservationQFunction{};
+      //    mass_conservation->AddDomainIntegrator(stokes_mass_conservation_qf, inputs,
+      //                                           outputs,
+      //                                           pressure_ir);
+      // }
+   }
+
+   void Mult(const Vector &x, Vector &r) const override
+   {
+      Vector xu(const_cast<double *>(x.Read()) + block_offsets[0],
+                block_offsets[1] - block_offsets[0]);
+      Vector xp(const_cast<double *>(x.Read()) + block_offsets[1],
+                block_offsets[2] - block_offsets[1]);
+      Vector ru(r.ReadWrite() + block_offsets[0],
+                block_offsets[1] - block_offsets[0]);
+      Vector rp(r.ReadWrite() + block_offsets[1],
+                block_offsets[2] - block_offsets[1]);
+
+      ParGridFunction p(&pressure_fes);
+      p.SetFromTrueDofs(xp);
+      momentum->SetParameters({&p, mesh_nodes});
+
+      momentum->Mult(xu, ru);
+      mass_conservation->Mult(xu, rp);
+
+      ru.SetSubVector(vel_ess_tdofs, 0.0);
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      jacobian_operator = std::make_shared<StokesJacobianOperator>(this, x);
+      return *jacobian_operator;
+
+      // fd_jacobian = std::make_shared<FDJacobian>(*this, x);
+      // return *fd_jacobian;
+   }
+
+   std::shared_ptr<DifferentiableOperator> momentum;
+   std::shared_ptr<DifferentiableOperator> continuity;
+
+   ParGridFunction *mesh_nodes;
+
+   ParMixedBilinearForm mass_conservation_form;
+   std::shared_ptr<Operator> mass_conservation;
+   const Array<int> block_offsets;
+   const Array<int> vel_ess_tdofs;
+
+   ParFiniteElementSpace &velocity_fes;
+   ParFiniteElementSpace &pressure_fes;
+
+   mutable std::shared_ptr<StokesJacobianOperator> jacobian_operator;
+   mutable std::shared_ptr<FDJacobian> fd_jacobian;
+};
+
+int main(int argc, char* argv[])
+{
+   constexpr int dim = 2;
+
+   Mpi::Init();
+
+   const char* device_config = "cpu";
+   const char* mesh_file = "../data/inline-quad.mesh";
+   int polynomial_order = 2;
+   int ir_order = 2;
+   int refinements = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.AddOption(&device_config, "-d", "--device",
+                  "Device configuration string, see Device::Configure().");
+   args.ParseCheck();
+
+   Device device(device_config);
+   if (Mpi::Root() == 0)
+   {
+      device.Print();
+   }
+
+   out << std::setprecision(8);
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.EnsureNodes();
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   H1_FECollection velocity_fec(polynomial_order, dim);
+   ParFiniteElementSpace velocity_fes(&mesh, &velocity_fec, dim);
+
+   H1_FECollection pressure_fec(polynomial_order - 1, dim);
+   ParFiniteElementSpace pressure_fes(&mesh, &pressure_fec);
+
+   out << velocity_fes.GetTrueVSize() << "\n";
+   out << pressure_fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule &velocity_ir =
+      IntRules.Get(velocity_fes.GetFE(0)->GetGeomType(),
+                   ir_order * velocity_fec.GetOrder());
+
+   const IntegrationRule &pressure_ir =
+      IntRules.Get(pressure_fes.GetFE(0)->GetGeomType(),
+                   ir_order * pressure_fec.GetOrder());
+
+   Array<int> bdr_attr_is_ess(mesh.bdr_attributes.Max());
+   bdr_attr_is_ess = 1;
+   Array<int> vel_ess_tdofs;
+   velocity_fes.GetEssentialTrueDofs(bdr_attr_is_ess, vel_ess_tdofs);
+
+   ParGridFunction u(&velocity_fes);
+   ParGridFunction p(&pressure_fes);
+
+   auto u_f = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      if (y >= 1.0)
+      {
+         u(0) = 1.0;
+      }
+      else
+      {
+         u(0) = 0.0;
+      }
+      u(1) = 0.0;
+   };
+   auto u_coef = VectorFunctionCoefficient(dim, u_f);
+
+   u.ProjectCoefficient(u_coef);
+   p = 0.0;
+
+   Array<int> block_offsets(3);
+   block_offsets[0] = 0;
+   block_offsets[1] = velocity_fes.GetTrueVSize();
+   block_offsets[2] = pressure_fes.GetTrueVSize();
+   block_offsets.PartialSum();
+
+   StokesOperator stokes(velocity_fes, pressure_fes, block_offsets,
+                         vel_ess_tdofs, velocity_ir, pressure_ir);
+
+   BlockVector x(block_offsets), y(block_offsets);
+   u.ParallelProject(x.GetBlock(0));
+   x.GetBlock(1) = 0.0;
+
+   GMRESSolver solver(MPI_COMM_WORLD);
+   solver.SetAbsTol(0.0);
+   solver.SetRelTol(1e-8);
+   // solver.SetKDim(100);
+   solver.SetMaxIter(500);
+   solver.SetPrintLevel(2);
+   // solver.SetPreconditioner(prec);
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.SetOperator(stokes);
+   newton.SetSolver(solver);
+   newton.SetRelTol(1e-6);
+   newton.SetMaxIter(50);
+   newton.SetPrintLevel(1);
+
+   Vector zero;
+   newton.Mult(zero, x);
+
+   u.SetFromTrueDofs(x.GetBlock(0));
+   p.SetFromTrueDofs(x.GetBlock(1));
+
+   ParaViewDataCollection dc("dfem_stokes", &mesh);
+   dc.SetHighOrderOutput(true);
+   dc.RegisterField("velocity", &u);
+   dc.RegisterField("pressure", &p);
+   dc.Save();
+
+   return 0;
+}

examples/dfem_test_diffusion.cpp

@@ -0,0 +1,193 @@
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+template <int dim = 2>
+int test_diffusion(
+   std::string mesh_file, int refinements, int polynomial_order)
+{
+   constexpr int num_samples = 1;
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction*>(mesh.GetNodes());
+   ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   L2_FECollection l2fec(0, dim);
+   ParFiniteElementSpace l2fes(&mesh, &l2fec);
+
+   out << "#dofs " << h1fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   printf("#ndof per el = %d\n", h1fes.GetFE(0)->GetDof());
+   printf("#nqp = %d\n", ir.GetNPoints());
+   printf("#q1d = %d\n", (int)floor(pow(ir.GetNPoints(), 1.0/dim) + 0.5));
+
+   std::shared_ptr<ParametricSpace> qdata_space;
+   if (mesh.GetElement(0)->GetType() == Element::QUADRILATERAL ||
+       mesh.GetElement(0)->GetType() == Element::HEXAHEDRON)
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            dim, dim * dim, ir.GetNPoints(), dim * dim * ir.GetNPoints() * mesh.GetNE());
+   }
+   else
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            1, dim * dim, ir.GetNPoints(), dim * dim * ir.GetNPoints() * mesh.GetNE());
+   }
+
+   ParametricFunction qdata(*qdata_space);
+
+   ParGridFunction f1_g(&h1fes);
+   ParGridFunction rho_g(&l2fes);
+
+   auto f1 = [](const Vector& coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      if (dim == 3)
+      {
+         const double z = coords(2);
+         return 2.345 + x + x*y + 1.25 * z*x;
+      }
+      else
+      {
+         return x + x*y + 2.345;
+      }
+   };
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   rho_g = 2.0;
+
+   Vector x(f1_g);
+
+   Vector y1(h1fes.GetTrueVSize());
+
+   auto diffusion_mf_kernel =
+      [] MFEM_HOST_DEVICE (
+         const tensor<real_t, dim>& dudxi,
+         const real_t& rho,
+         const tensor<real_t, dim, dim>& J,
+         const real_t& w)
+   {
+      auto invJ = inv(J);
+      return mfem::tuple{(pow(rho, 3.0)*(dudxi * invJ)) * transpose(invJ) * det(J) * w};
+   };
+
+   constexpr int Potential = 3;
+   constexpr int Diffusivity = 44;
+   constexpr int Coordinates = 55;
+
+   auto input_operators = mfem::tuple
+   {
+      Gradient<Potential>{},
+      Value<Diffusivity>{},
+      Gradient<Coordinates>{},
+      Weight{}
+   };
+   auto output_operator = mfem::tuple{Gradient<Potential>{}};
+
+   auto solutions = std::vector
+   {
+      FieldDescriptor{Potential, &h1fes}
+   };
+   auto parameters = std::vector
+   {
+      FieldDescriptor{Diffusivity, &l2fes},
+      FieldDescriptor{Coordinates, &mesh_fes}
+   };
+
+   DifferentiableOperator dop(solutions, parameters, mesh);
+   auto derivatives = std::integer_sequence<size_t, Potential, Diffusivity> {};
+   Array<int> domain_attributes(mesh.attributes.Size());
+   domain_attributes = 1;
+   dop.AddDomainIntegrator(
+      diffusion_mf_kernel, input_operators, output_operator, ir, domain_attributes,
+      derivatives);
+
+   dop.SetParameters({&rho_g, mesh_nodes});
+   StopWatch sw;
+   sw.Start();
+   for (int i = 0; i < num_samples; i++)
+   {
+      dop.Mult(x, y1);
+   }
+   sw.Stop();
+   printf("dfem mf:       %fs\n", sw.RealTime() / num_samples);
+   y1.HostRead();
+
+   auto dfdp = dop.GetDerivative(Diffusivity, {&f1_g}, {&rho_g, mesh_nodes});
+
+   dfdp->Mult(rho_g, y1);
+
+   // printf("y1: ");
+   // print_vector(y1);
+
+   {
+      // Create a direction vector for rho
+      Vector dir(rho_g);
+
+      // Small parameter for finite difference
+      double eps = 1.0e-6;
+
+      // Compute f(rho + eps*dir)
+      Vector rho_plus(rho_g);
+      rho_plus.Add(eps, dir);
+      dop.SetParameters({&rho_plus, mesh_nodes});
+      Vector f_plus(x.Size());
+      dop.Mult(x, f_plus);
+
+      // Compute f(rho - eps*dir)
+      Vector rho_minus(rho_g);
+      rho_minus.Add(-eps, dir);
+      dop.SetParameters({&rho_minus, mesh_nodes});
+      Vector f_minus(x.Size());
+      dop.Mult(x, f_minus);
+
+      // Finite difference approximation of the derivative action
+      Vector fd_result(x.Size());
+      subtract(f_plus, f_minus, fd_result);
+      fd_result *= 1.0/(2.0*eps);
+
+      // printf("fd: ");
+      // print_vector(fd_result);
+
+      fd_result -= y1;
+      double absolute_error = fd_result.Norml2();
+      double relative_error = absolute_error / y1.Norml2();
+      out << "Absolute error ||dFdrho_FD * rho - dfem||_l2 = " << absolute_error <<
+          "\n";
+      out << "Relative error ||dFdrho_FD * rho - dfem||_l2 / ||dfem||_l2 = " <<
+          relative_error << "\n";      // if (frhopv.Norml2() > eps)
+      // {
+      //    out << "||dFdu_FD u^* - ex||_l2 = " << frhopv.Norml2() << "\n";
+      //    return 1;
+      // }
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_diffusion<2>);

examples/dfem_test_diffusion_2d.cpp

@@ -0,0 +1,174 @@
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_diffusion(
+   std::string mesh_file, int refinements, int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == 2, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction*>(mesh.GetNodes());
+   ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   out << "#dofs " << h1fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder());
+
+   out << "#qp: " << ir.GetNPoints() << "\n";
+
+   ParGridFunction f1_g(&h1fes);
+   ParGridFunction rho_g(&h1fes);
+
+   constexpr int Potential = 0;
+   constexpr int Coordinates = 1;
+
+   auto kernel = [] MFEM_HOST_DEVICE(
+                    const tensor<double, 2, 2>& J,
+                    const double& w, const tensor<double, 2>& dudxi)
+   {
+      auto invJ = inv(J);
+      return mfem::tuple{dudxi * invJ * transpose(invJ) * det(J) * w};
+   };
+
+   mfem::tuple argument_operators =
+   {
+      Gradient<Coordinates>{}, Weight{}, Gradient<Potential>{}
+   };
+   mfem::tuple output_operator = {Gradient<Potential>{}};
+
+   auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+   auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes}};
+
+   DifferentiableOperator dop(solutions, parameters, mesh);
+
+   auto f1 = [](const Vector& coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      return 2.345 + 0.25 * x * x * y + y * y * x;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(f1_g), y(h1fes.TrueVSize());
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+   y.HostRead();
+
+   ParBilinearForm a(&h1fes);
+   a.AddDomainIntegrator(new DiffusionIntegrator);
+   a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   a.Assemble();
+   a.Finalize();
+
+   Vector y2(h1fes.TrueVSize());
+   a.Mult(x, y2);
+   y2.HostRead();
+   Vector diff(y2);
+   diff -= y;
+   if (diff.Norml2() > 1e-10)
+   {
+      print_vector(diff);
+      print_vector(y2);
+      print_vector(y);
+      return 1;
+   }
+
+   // // Test linearization here as well
+   // auto dFdu = dop.GetDerivativeWrt<0>({&f1_g}, {mesh_nodes});
+
+   // if (dFdu->Height() != h1fes.GetTrueVSize())
+   // {
+   //    out << "dFdu unexpected height of " << dFdu->Height() << "\n";
+   //    return 1;
+   // }
+
+   // dFdu->Mult(x, y);
+   // y.HostRead();
+   // a.Mult(x, y2);
+   // y2.HostRead();
+
+   // diff = y2;
+   // diff -= y;
+   // if (diff.Norml2() > 1e-10)
+   // {
+   //    print_vector(diff);
+   //    print_vector(y2);
+   //    print_vector(y);
+   //    return 1;
+   // }
+
+   // // fd jacobian test
+   // {
+   //    double eps = 1.0e-6;
+   //    Vector v(x), xpv(x), xmv(x), fxpv(x.Size()), fxmv(x.Size());
+   //    v *= eps;
+   //    xpv += v;
+   //    xmv -= v;
+   //    dop.Mult(xpv, fxpv);
+   //    dop.Mult(xmv, fxmv);
+   //    fxpv -= fxmv;
+   //    fxpv /= (2.0*eps);
+
+   //    fxpv -= y;
+   //    if (fxpv.Norml2() > eps)
+   //    {
+   //       out << "||dFdu_FD u^* - ex||_l2 = " << fxpv.Norml2() << "\n";
+   //       return 1;
+   //    }
+   // }
+
+   // f1_g.ProjectCoefficient(f1_c);
+   // rho_g.ProjectCoefficient(rho_c);
+   // auto dFdrho = dop.GetDerivativeWrt<1>({&f1_g}, {&rho_g, mesh_nodes});
+   // if (dFdrho->Height() != h1fes.GetTrueVSize())
+   // {
+   //    out << "dFdrho unexpected height of " << dFdrho->Height() << "\n";
+   //    return 1;
+   // }
+
+   // dFdrho->Mult(rho_g, y);
+
+   // // fd test
+   // {
+   //    double eps = 1.0e-6;
+   //    Vector v(rho_g), rhopv(rho_g), rhomv(rho_g), frhopv(x.Size()),
+   //    frhomv(x.Size()); v *= eps; rhopv += v; rhomv -= v;
+   //    dop.SetParameters({&rhopv, mesh_nodes});
+   //    dop.Mult(x, frhopv);
+   //    dop.SetParameters({&rhomv, mesh_nodes});
+   //    dop.Mult(x, frhomv);
+   //    frhopv -= frhomv;
+   //    frhopv /= (2.0*eps);
+
+   //    frhopv -= y;
+   //    if (frhopv.Norml2() > eps)
+   //    {
+   //       out << "||dFdu_FD u^* - ex||_l2 = " << frhopv.Norml2() << "\n";
+   //       return 1;
+   //    }
+   // }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_diffusion);

examples/dfem_test_diffusion_3d.cpp

@@ -0,0 +1,370 @@
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_diffusion_3d(
+   std::string mesh_file, int refinements, int polynomial_order)
+{
+   constexpr int num_samples = 1;
+   constexpr int dim = 2;
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "incorrect mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction*>(mesh.GetNodes());
+   ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   out << "#dofs " << h1fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   printf("#ndof per el = %d\n", h1fes.GetFE(0)->GetDof());
+   printf("#nqp = %d\n", ir.GetNPoints());
+   printf("#q1d = %d\n", (int)floor(pow(ir.GetNPoints(), 1.0/dim) + 0.5));
+
+   std::shared_ptr<ParametricSpace> qdata_space;
+   if (mesh.GetElement(0)->GetType() == Element::QUADRILATERAL ||
+       mesh.GetElement(0)->GetType() == Element::HEXAHEDRON)
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            dim, dim * dim, ir.GetNPoints(), dim * dim * ir.GetNPoints() * mesh.GetNE());
+   }
+   else
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            1, dim * dim, ir.GetNPoints(), dim * dim * ir.GetNPoints() * mesh.GetNE());
+   }
+
+   ParametricFunction qdata(*qdata_space);
+
+   ParGridFunction f1_g(&h1fes);
+   ParGridFunction rho_g(&h1fes);
+
+   auto f1 = [](const Vector& coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      if (dim == 3)
+      {
+         const double z = coords(2);
+         return 2.345 + x + x*y + 1.25 * z*x;
+      }
+      else
+      {
+         return x;
+      }
+   };
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   std::unique_ptr<DerivativeOperator> dfdu;
+
+   Vector x(f1_g), y(h1fes.GetTrueVSize());
+   {
+      auto diffusion_mf_kernel =
+         [] MFEM_HOST_DEVICE (
+            const tensor<real_t, dim>& dudxi,
+            const tensor<real_t, dim, dim>& J,
+            const real_t& w)
+      {
+         auto invJ = inv(J);
+         return mfem::tuple{((dudxi * invJ)) * transpose(invJ) * det(J) * w};
+      };
+
+      constexpr int Potential = 0;
+      constexpr int Coordinates = 1;
+
+      auto input_operators = mfem::tuple{Gradient<Potential>{}, Gradient<Coordinates>{}, Weight{}};
+      auto output_operator = mfem::tuple{Gradient<Potential>{}};
+
+      auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+      auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes}};
+
+      DifferentiableOperator dop(solutions, parameters, mesh);
+      auto derivatives = std::integer_sequence<size_t, Potential> {};
+      dop.AddDomainIntegrator(
+         diffusion_mf_kernel, input_operators, output_operator, ir, derivatives);
+
+      dop.SetParameters({mesh_nodes});
+      StopWatch sw;
+      sw.Start();
+      for (int i = 0; i < num_samples; i++)
+      {
+         dop.Mult(x, y);
+      }
+      sw.Stop();
+      printf("dfem mf:       %fs\n", sw.RealTime() / num_samples);
+      y.HostRead();
+
+      dfdu = dop.GetDerivative(Potential, {&f1_g}, {mesh_nodes});
+
+      // sw.Start();
+      // for (int i = 0; i < num_samples; i++)
+      // {
+      dfdu->Mult(x, y);
+      // }
+      // sw.Stop();
+      // printf("dfem mf VJP:       %fs\n", sw.RealTime() / num_samples);
+
+      printf("y: ");
+      print_vector(y);
+   }
+
+   // {
+   //    auto diffusion_setup_kernel =
+   //       [] MFEM_HOST_DEVICE (
+   //          const tensor<double, dim, dim>& J,
+   //          const double& w)
+   //    {
+   //       auto invJ = inv(J);
+   //       tensor<real_t, dim, dim> C{0.0};
+   //       C(0, 0) = 2.0;
+   //       C(1, 0) = 3.0;
+   //       C(1, 1) = 4.0;
+   //       if (dim == 3)
+   //       {
+   //          C(2, 2) = 1.0;
+   //       }
+   //       return mfem::tuple{C * invJ * transpose(invJ) * det(J) * w};
+   //    };
+
+   //    constexpr int Potential = 0;
+   //    constexpr int Coordinates = 1;
+   //    constexpr int QData = 2;
+
+   //    auto input_operators = mfem::tuple{Gradient<Coordinates>{}, Weight{}};
+   //    auto output_operator = mfem::tuple{None<QData>{}};
+
+   //    auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+   //    auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes},
+   //                                  FieldDescriptor{QData, qdata_space.get()}};
+
+   //    DifferentiableOperator dop(solutions, parameters, mesh);
+   //    dop.AddDomainIntegrator(
+   //       diffusion_setup_kernel, input_operators, output_operator, ir);
+
+   //    dop.SetParameters({mesh_nodes, &qdata});
+   //    StopWatch sw;
+   //    sw.Start();
+   //    for (int i = 0; i < num_samples; i++)
+   //    {
+   //       dop.Mult(x, qdata);
+   //    }
+   //    sw.Stop();
+   //    printf("dfem pa setup: %fs\n", sw.RealTime() / num_samples);
+   //    qdata.HostRead();
+   // }
+
+   // printf("qdata: ");
+   // print_vector(qdata);
+
+   // {
+   //    auto diffusion_apply_kernel =
+   //       [] MFEM_HOST_DEVICE (
+   //          const tensor<real_t, dim>& dudxi,
+   //          const tensor<double, dim, dim>& qdata)
+   //    {
+   //       return mfem::tuple{qdata * dudxi};
+   //    };
+
+   //    constexpr int Potential = 0;
+   //    constexpr int QData = 1;
+
+   //    auto input_operators = mfem::tuple{Gradient<Potential>{}, None<QData>{}};
+   //    auto output_operator = mfem::tuple{Gradient<Potential>{}};
+
+   //    auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+   //    auto parameters = std::vector{FieldDescriptor{QData, qdata_space.get()}};
+
+   //    DifferentiableOperator dop(solutions, parameters, mesh);
+   //    dop.AddDomainIntegrator(
+   //       diffusion_apply_kernel, input_operators, output_operator, ir);
+
+   //    dop.SetParameters({&qdata});
+   //    StopWatch sw;
+   //    sw.Start();
+   //    for (int i = 0; i < num_samples; i++)
+   //    {
+   //       dop.Mult(x, y);
+   //    }
+   //    sw.Stop();
+   //    printf("dfem pa apply: %fs\n", sw.RealTime() / num_samples);
+   //    y.HostRead();
+   // }
+
+   // printf("y: ");
+   // print_vector(y);
+
+   Vector y2(h1fes.TrueVSize());
+   {
+      ParBilinearForm a(&h1fes);
+      auto diff_integ = new DiffusionIntegrator;
+      diff_integ->SetIntRule(&ir);
+      a.AddDomainIntegrator(diff_integ);
+      if (mesh.GetElement(0)->GetType() == Element::QUADRILATERAL ||
+          mesh.GetElement(0)->GetType() == Element::HEXAHEDRON)
+      {
+         // a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+      }
+
+      OperatorPtr A;
+      StopWatch sw;
+      sw.Start();
+      a.Assemble();
+      a.Finalize();
+      Array<int> empty;
+      a.FormSystemMatrix(empty, A);
+      sw.Stop();
+      printf("mfem pa setup: %fs\n", sw.RealTime());
+
+      sw.Clear();
+      sw.Start();
+      y2 = 0.0;
+      for (int i = 0; i < num_samples; i++)
+      {
+         A->Mult(x, y2);
+      }
+      sw.Stop();
+      printf("mfem pa apply: %fs\n", sw.RealTime() / num_samples);
+      y2.HostRead();
+   }
+
+   printf("y2: ");
+   print_vector(y2);
+
+   Vector diff(y2);
+   diff -= y;
+   if (diff.Norml2() > 1e-12)
+   {
+      // printf("y: ");
+      // print_vector(y);
+      // printf("y2: ");
+      // print_vector(y2);
+      // printf("diff: ");
+      // print_vector(diff);
+      return 1;
+   }
+
+   // {
+   //    DenseMatrix m(dim);
+   //    m(0, 0) = 2.0;
+   //    m(1, 0) = 3.0;
+   //    m(1, 1) = 4.0;
+   //    if (dim == 3)
+   //    {
+   //       m(2, 2) = 1.0;
+   //    }
+   //    MatrixConstantCoefficient matrix_coeff(m);
+   //    ParBilinearForm a(&h1fes);
+   //    auto diff_integ = new DiffusionIntegrator(matrix_coeff);
+   //    diff_integ->SetIntRule(&ir);
+   //    a.AddDomainIntegrator(diff_integ);
+
+   //    OperatorPtr A;
+   //    a.Assemble();
+   //    a.Finalize();
+   //    Array<int> empty;
+   //    a.FormSystemMatrix(empty, A);
+   //    A->MultTranspose(x, y);
+   //    // out << "mfem A^T * x transpose\n";
+   //    // print_vector(y);
+   // }
+
+   // Test linearization here as well
+   // auto dFdu = dop.GetDerivativeWrt<0>({&f1_g}, {mesh_nodes});
+
+   // if (dFdu->Height() != h1fes.GetTrueVSize())
+   // {
+   //    out << "dFdu unexpected height of " << dFdu->Height() << "\n";
+   //    return 1;
+   // }
+
+   // dFdu->Mult(x, y);
+   // y.HostRead();
+   // a.Mult(x, y2);
+   // y2.HostRead();
+
+   // diff = y2;
+   // diff -= y;
+   // if (diff.Norml2() > 1e-10)
+   // {
+   //    print_vector(diff);
+   //    print_vector(y2);
+   //    print_vector(y);
+   //    return 1;
+   // }
+
+   // // fd jacobian test
+   // {
+   //    double eps = 1.0e-6;
+   //    Vector v(x), xpv(x), xmv(x), fxpv(x.Size()), fxmv(x.Size());
+   //    v *= eps;
+   //    xpv += v;
+   //    xmv -= v;
+   //    dpotential->Mult(xpv, fxpv);
+   //    dpotential->Mult(xmv, fxmv);
+   //    fxpv -= fxmv;
+   //    fxpv /= (2.0*eps);
+
+   //    fxpv -= y;
+   //    if (fxpv.Norml2() > eps)
+   //    {
+   //       out << "||dFdu_FD u^* - ex||_l2 = " << fxpv.Norml2() << "\n";
+   //       return 1;
+   //    }
+   // }
+
+   // f1_g.ProjectCoefficient(f1_c);
+   // rho_g.ProjectCoefficient(rho_c);
+   // auto dFdrho = dop.GetDerivativeWrt<1>({&f1_g}, {&rho_g, mesh_nodes});
+   // if (dFdrho->Height() != h1fes.GetTrueVSize())
+   // {
+   //    out << "dFdrho unexpected height of " << dFdrho->Height() << "\n";
+   //    return 1;
+   // }
+
+   // dFdrho->Mult(rho_g, y);
+
+   // // fd test
+   // {
+   //    double eps = 1.0e-6;
+   //    Vector v(rho_g), rhopv(rho_g), rhomv(rho_g), frhopv(x.Size()),
+   //    frhomv(x.Size()); v *= eps; rhopv += v; rhomv -= v;
+   //    dop.SetParameters({&rhopv, mesh_nodes});
+   //    dop.Mult(x, frhopv);
+   //    dop.SetParameters({&rhomv, mesh_nodes});
+   //    dop.Mult(x, frhomv);
+   //    frhopv -= frhomv;
+   //    frhopv /= (2.0*eps);
+
+   //    frhopv -= y;
+   //    if (frhopv.Norml2() > eps)
+   //    {
+   //       out << "||dFdu_FD u^* - ex||_l2 = " << frhopv.Norml2() << "\n";
+   //       return 1;
+   //    }
+   // }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_diffusion_3d);

examples/dfem_test_elasticity.cpp

@@ -0,0 +1,109 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_elasticity(std::string mesh_file,
+                    int refinements,
+                    int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   const int vdim = dim;
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   Array<int> ess_tdof;
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 6 * h1fec.GetOrder());
+
+   out << "#qp: " << ir.GetNPoints() << "\n";
+   out << "#dof_el: " << h1fes.GetRestrictionMatrix()->Height() / mesh.GetNE() <<
+       "\n";
+
+   ParGridFunction u(&h1fes);
+
+   auto f1 = [](const Vector& coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      u(0) = 2.345 + 0.25 * x * x * y + y * y * x;
+      u(1) = 2.345 - 0.25 * x * y * y + y * x * x;
+   };
+
+   VectorFunctionCoefficient u_c(dim, f1);
+   u.ProjectCoefficient(u_c);
+
+   ConstantCoefficient l_coeff(0.5), m_coeff(0.25);
+
+   ParBilinearForm A_form(&h1fes);
+   auto A_integ = new ElasticityIntegrator(l_coeff, m_coeff);
+   A_integ->SetIntegrationRule(ir);
+   A_form.AddDomainIntegrator(A_integ);
+   A_form.Assemble();
+   A_form.Finalize();
+
+   auto elasticity_kernel = [](const tensor<double, 2, 2> &dudxi,
+                               const tensor<double, 2, 2> &J,
+                               const double &w)
+   {
+      constexpr double lambda = 0.5;
+      constexpr double mu = 0.25;
+      static constexpr auto I = mfem::internal::IsotropicIdentity<2>();
+      auto invJ = inv(J);
+      auto eps = sym(dudxi * invJ);
+      return mfem::tuple{transpose(lambda * tr(eps) * I + 2.0 * mu * eps) * det(J) * w * transpose(invJ)};
+   };
+
+   mfem::tuple argument_operators{Gradient{"displacement"}, Gradient{"coordinates"}, Weight{}};
+   mfem::tuple output_operator{Gradient{"displacement"}};
+
+   ElementOperator op{elasticity_kernel, argument_operators, output_operator};
+
+   std::array solutions{FieldDescriptor{&h1fes, "displacement"}};
+   std::array parameters{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop{solutions, parameters, mfem::tuple{op}, mesh, ir};
+
+   Vector x(u), y1(h1fes.GetTrueVSize()),
+          y2(h1fes.GetTrueVSize());
+
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y1);
+   y1.HostRead();
+
+   A_form.Mult(x, y2);
+   y2.HostRead();
+
+   Vector diff(y2);
+   diff -= y1;
+   if (diff.Norml2() > 1e-10)
+   {
+      out << "||F(u) - ex||_l2 = " << diff.Norml2() << "\n";
+      print_vector(diff);
+      print_vector(y1);
+      print_vector(y2);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_elasticity);

examples/dfem_test_interpolate_gradient_linear_scalar.cpp

@@ -0,0 +1,148 @@
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_interpolate_gradient_linear_scalar(std::string mesh_file,
+                                            int refinements,
+                                            int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   const int dim = mesh_serial.Dimension();
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   ParGridFunction *mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   ParGridFunction f1_g(&h1fes);
+
+   std::shared_ptr<ParametricSpace> qdata_space;
+   if (mesh.GetElement(0)->GetType() == Element::QUADRILATERAL ||
+       mesh.GetElement(0)->GetType() == Element::HEXAHEDRON)
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            dim, dim, ir.GetNPoints(), dim * ir.GetNPoints() * mesh.GetNE());
+   }
+   else
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            1, dim, ir.GetNPoints(), dim * ir.GetNPoints() * mesh.GetNE());
+   }
+
+   ParametricFunction qdata(*qdata_space);
+
+   auto kernel_2d = [] MFEM_HOST_DEVICE (
+                       const tensor<real_t, 2> &dudxi,
+                       const tensor<real_t, 2, 2> &J)
+   {
+      out << "J: " << J << std::endl;
+      out << "dudxi: " << dudxi << std::endl;
+      out << inv(J) << std::endl;
+      return mfem::tuple{dudxi * inv(J)};
+   };
+
+   auto kernel_3d = [] MFEM_HOST_DEVICE (
+                       const tensor<real_t, 3> &dudxi,
+                       const tensor<real_t, 3, 3> &J)
+   {
+      out << "J: " << J << std::endl;
+      out << "dudxi: " << dudxi << std::endl;
+      out << inv(J) << std::endl;
+      return mfem::tuple{dudxi * inv(J)};
+   };
+
+   constexpr int Potential = 0;
+   constexpr int Coordinates = 1;
+   constexpr int Qdata = 2;
+
+   auto input_operators = mfem::tuple{Gradient<Potential>{}, Gradient<Coordinates>{}};
+   auto output_operator = mfem::tuple{None<Qdata>{}};
+
+   auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+   auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes},
+                                 FieldDescriptor{Qdata, qdata_space.get()}};
+
+   DifferentiableOperator dop(solutions, parameters, mesh);
+   if (dim == 2)
+   {
+      dop.AddDomainIntegrator(kernel_2d, input_operators, output_operator, ir);
+   }
+   else
+   {
+      dop.AddDomainIntegrator(kernel_3d, input_operators, output_operator, ir);
+   }
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      if (coords.Size() > 2)
+      {
+         const double z = coords(2);
+         return 2.345 + x * y * z + y * z;
+      }
+      else
+      {
+         return 2.345 + x * y + y;
+      }
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(*f1_g.GetTrueDofs());
+   dop.SetParameters({mesh_nodes, &qdata});
+   dop.Mult(x, qdata);
+
+   Vector f_test(qdata.Size());
+   for (int e = 0; e < mesh.GetNE(); e++)
+   {
+      ElementTransformation *T = mesh.GetElementTransformation(e);
+
+      for (int qp = 0; qp < ir.GetNPoints(); qp++)
+      {
+         const IntegrationPoint &ip = ir.IntPoint(qp);
+         T->SetIntPoint(&ip);
+
+         Vector g(dim);
+         f1_g.GetGradient(*T, g);
+         // printf("(%f, %f, %f): (%f, %f, %f)\n", ip.x, ip.y, ip.z,  g(0), g(1), g(2));
+         for (int d = 0; d < dim; d++)
+         {
+            int qpo = qp * dim;
+            int eo = e * (ir.GetNPoints() * dim);
+            f_test(d + qpo + eo) = g(d);
+         }
+      }
+   }
+
+   Vector diff(f_test);
+   diff -= qdata;
+   if (diff.Norml2() > 1e-12)
+   {
+      print_vector(diff);
+      print_vector(f_test);
+      print_vector(qdata);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_interpolate_gradient_linear_scalar);

examples/dfem_test_interpolate_gradient_linear_scalar_3d.cpp

@@ -0,0 +1,113 @@
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_interpolate_gradient_linear_scalar_3d(std::string mesh_file,
+                                               int refinements,
+                                               int polynomial_order)
+{
+   constexpr int dim = 3;
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "wrong mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   mesh_serial.Clear();
+
+   ParGridFunction *mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   // const IntegrationRule &ir =
+   //    IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+   IntegrationRules gll_rules(0, Quadrature1D::GaussLobatto);
+   const IntegrationRule &ir = gll_rules.Get(h1fes.GetFE(0)->GetGeomType(),
+                                             2 * polynomial_order - 1);
+
+   ParGridFunction f1_g(&h1fes);
+
+   ParametricSpace pspace(dim, dim, ir.GetNPoints(),
+                          dim * ir.GetNPoints() * mesh.GetNE());
+   ParametricFunction qdata(pspace);
+
+   auto kernel = [](const tensor<double, dim> &dudxi,
+                    const tensor<double, dim, dim> &J)
+   {
+      return mfem::tuple{dudxi * inv(J)};
+   };
+
+   mfem::tuple argument_operators = {Gradient{"potential"}, Gradient{"coordinates"}};
+   mfem::tuple output_operator = {None{"qdata"}};
+
+   ElementOperator eop = {kernel, argument_operators, output_operator};
+   auto ops = mfem::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "potential"}};
+   auto parameters = std::array
+   {
+      FieldDescriptor{&mesh_fes, "coordinates"},
+      FieldDescriptor{&pspace, "qdata"}
+   };
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      const double z = coords(2);
+      return 2.345 + x * y * z + y * z;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(*f1_g.GetTrueDofs()), y(h1fes.TrueVSize() * dim);
+   dop.SetParameters({mesh_nodes, &qdata});
+   dop.Mult(x, y);
+
+   Vector f_test(h1fes.GetElementRestriction(
+                    ElementDofOrdering::LEXICOGRAPHIC)->Height() * dim);
+   for (int e = 0; e < mesh.GetNE(); e++)
+   {
+      ElementTransformation *T = mesh.GetElementTransformation(e);
+
+      for (int qp = 0; qp < ir.GetNPoints(); qp++)
+      {
+         const IntegrationPoint &ip = ir.IntPoint(qp);
+         T->SetIntPoint(&ip);
+
+         Vector g(dim);
+         f1_g.GetGradient(*T, g);
+         // printf("(%f, %f, %f): (%f, %f, %f)\n", ip.x, ip.y, ip.z,  g(0), g(1), g(2));
+         for (int d = 0; d < dim; d++)
+         {
+            int qpo = qp * dim;
+            int eo = e * (ir.GetNPoints() * dim);
+            f_test(d + qpo + eo) = g(d);
+         }
+      }
+   }
+
+   Vector diff(f_test);
+   diff -= y;
+   if (diff.Norml2() > 1e-10)
+   {
+      print_vector(diff);
+      print_vector(f_test);
+      print_vector(y);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_interpolate_gradient_linear_scalar_3d);

examples/dfem_test_interpolate_linear_scalar.cpp

@@ -0,0 +1,108 @@
+#include "dfem/dfem_test_macro.hpp"
+#include "examples/dfem/dfem_parametricspace.hpp"
+#include "examples/dfem/dfem_util.hpp"
+#include "fem/geom.hpp"
+#include <memory>
+
+using namespace mfem;
+
+int test_interpolate_linear_scalar(std::string mesh_file,
+                                   int refinements,
+                                   int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   ParGridFunction f1_g(&h1fes);
+
+   std::shared_ptr<ParametricSpace> qdata_space;
+   if (mesh.GetElement(0)->GetType() == Element::QUADRILATERAL ||
+       mesh.GetElement(0)->GetType() == Element::HEXAHEDRON)
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            dim, 1, ir.GetNPoints(), ir.GetNPoints() * mesh.GetNE());
+   }
+   else
+   {
+      qdata_space =
+         std::make_shared<ParametricSpace>(
+            1, 1, ir.GetNPoints(), ir.GetNPoints() * mesh.GetNE());
+   }
+
+   ParametricFunction qdata(*qdata_space);
+
+   auto kernel = [] MFEM_HOST_DEVICE (const double &u)
+   {
+      return mfem::tuple{u};
+   };
+
+   constexpr int Potential = 0;
+   constexpr int Qdata = 1;
+
+   auto input_operators = mfem::tuple {Value<Potential> {}};
+   auto output_operator = mfem::tuple {None<Qdata> {}};
+
+   auto solutions = std::vector {FieldDescriptor{Potential, &h1fes}};
+   auto parameters = std::vector {FieldDescriptor{Qdata, qdata_space.get()}};
+
+   DifferentiableOperator dop(solutions, parameters, mesh);
+   dop.AddDomainIntegrator(kernel, input_operators, output_operator, ir);
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      return 2.345 + x + y;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(*f1_g.GetTrueDofs());
+   dop.SetParameters({&qdata});
+   dop.Mult(x, qdata);
+
+   Vector f_test(ir.GetNPoints() * mesh.GetNE());
+   for (int e = 0; e < mesh.GetNE(); e++)
+   {
+      ElementTransformation *T = mesh.GetElementTransformation(e);
+      for (int qp = 0; qp < ir.GetNPoints(); qp++)
+      {
+         const IntegrationPoint &ip = ir.IntPoint(qp);
+         T->SetIntPoint(&ip);
+         f_test((e * ir.GetNPoints()) + qp) = f1_c.Eval(*T, ip);
+      }
+   }
+
+   Vector diff(f_test);
+   diff -= qdata;
+   if (diff.Norml2() > 1e-12)
+   {
+      print_vector(diff);
+      print_vector(f_test);
+      print_vector(qdata);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_interpolate_linear_scalar);

examples/dfem_test_interpolate_linear_scalar_3d.cpp

@@ -0,0 +1,93 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_interpolate_linear_scalar_3d(std::string mesh_file,
+                                      int refinements,
+                                      int polynomial_order)
+{
+   constexpr int dim = 3;
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "wrong mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   mesh_serial.Clear();
+
+   ParGridFunction *mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   ParGridFunction f1_g(&h1fes);
+
+   auto kernel = [](const double &u)
+   {
+      return mfem::tuple{u};
+   };
+
+   mfem::tuple argument_operators = {Value{"potential"}};
+   mfem::tuple output_operator = {None{"potential"}};
+
+   ElementOperator eop = {kernel, argument_operators, output_operator};
+   auto ops = mfem::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "potential"}};
+   auto parameters = std::array{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      const double z = coords(2);
+      return 2.345 + x + y + 1.25 * z;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(*f1_g.GetTrueDofs()), y(h1fes.TrueVSize());
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+
+   Vector f_test(h1fes.GetElementRestriction(
+                    ElementDofOrdering::LEXICOGRAPHIC)->Height());
+   for (int e = 0; e < mesh.GetNE(); e++)
+   {
+      ElementTransformation *T = mesh.GetElementTransformation(e);
+      for (int qp = 0; qp < ir.GetNPoints(); qp++)
+      {
+         const IntegrationPoint &ip = ir.IntPoint(qp);
+         T->SetIntPoint(&ip);
+
+         f_test((e * ir.GetNPoints()) + qp) = f1_c.Eval(*T, ip);
+      }
+   }
+
+   Vector diff(f_test);
+   diff -= y;
+   if (diff.Norml2() > 1e-10)
+   {
+      print_vector(diff);
+      print_vector(f_test);
+      print_vector(y);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_interpolate_linear_scalar_3d);

examples/dfem_test_interpolate_linear_vector.cpp

@@ -0,0 +1,104 @@
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_interpolate_linear_vector(std::string mesh_file, int refinements,
+                                   int polynomial_order)
+{
+   constexpr int vdim = 2;
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   ParGridFunction f1_g(&h1fes);
+
+   // ParametricSpace qdata_space(dim, 1, ir.GetNPoints(),
+   //                             ir.GetNPoints() * mesh.GetNE());
+   ParametricSpace qdata_space(1, vdim, ir.GetNPoints(),
+                               vdim * ir.GetNPoints() * mesh.GetNE());
+
+   ParametricFunction qdata(qdata_space);
+
+   auto kernel = [] MFEM_HOST_DEVICE (const tensor<real_t, vdim> &u)
+   {
+      return mfem::tuple{u};
+   };
+
+   constexpr int Potential = 0;
+   constexpr int Qdata = 1;
+
+   auto input_operators = mfem::tuple{Value<Potential>{}};
+   auto output_operator = mfem::tuple{None<Qdata>{}};
+
+   auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+   auto parameters = std::vector{FieldDescriptor{Qdata, &qdata_space}};
+
+   DifferentiableOperator dop(solutions, parameters, mesh);
+   dop.AddDomainIntegrator(kernel, input_operators, output_operator, ir);
+
+   auto f1 = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      u(0) = 2.345 + x + y;
+      u(1) = 12.345 + x + y;
+   };
+
+   VectorFunctionCoefficient f1_c(vdim, f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(*f1_g.GetTrueDofs());
+   dop.SetParameters({&qdata});
+   dop.Mult(x, qdata);
+
+   Vector f_test(qdata.Size());
+   for (int e = 0; e < mesh.GetNE(); e++)
+   {
+      ElementTransformation *T = mesh.GetElementTransformation(e);
+      for (int qp = 0; qp < ir.GetNPoints(); qp++)
+      {
+         const IntegrationPoint &ip = ir.IntPoint(qp);
+         T->SetIntPoint(&ip);
+
+         Vector f(vdim);
+         f1_g.GetVectorValue(*T, ip, f);
+         for (int d = 0; d < vdim; d++)
+         {
+            int qpo = qp * vdim;
+            int eo = e * (ir.GetNPoints() * vdim);
+            f_test(d + qpo + eo) = f(d);
+         }
+      }
+   }
+
+   Vector diff(f_test);
+   diff -= qdata;
+   if (diff.Norml2() > 1e-12)
+   {
+      print_vector(diff);
+      print_vector(f_test);
+      print_vector(qdata);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_interpolate_linear_vector);

examples/dfem_test_interpolate_linear_vector_3d.cpp

@@ -0,0 +1,105 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_interpolate_linear_vector_3d(std::string mesh_file, int refinements,
+                                      int polynomial_order)
+{
+
+   constexpr int dim = 3;
+   constexpr int vdim = 3;
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "wrong mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+   mesh.SetCurvature(1);
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   QuadratureSpace qspace(mesh, ir);
+   QuadratureFunction qf(&qspace, vdim);
+
+   ParGridFunction f1_g(&h1fes);
+
+   auto kernel = [](const tensor<double, vdim> &u)
+   {
+      return mfem::tuple{u};
+   };
+
+   mfem::tuple argument_operators = {Value{"potential"}};
+   mfem::tuple output_operator = {None{"potential"}};
+
+   ElementOperator eop{kernel, argument_operators, output_operator};
+   auto ops = mfem::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "potential"}};
+   auto parameters = std::array{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   auto f1 = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      const double z = coords(2);
+      u(0) = 2.345 + x + y + 3.0 * z;
+      u(1) = 12.345 + x + y + 2.0 * z;
+      u(2) = 5.345 + x + y + 1.0 * z;
+   };
+
+   VectorFunctionCoefficient f1_c(vdim, f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(f1_g), y(f1_g.Size());
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+
+   Vector f_test(qf.Size());
+   for (int e = 0; e < mesh.GetNE(); e++)
+   {
+      ElementTransformation *T = mesh.GetElementTransformation(e);
+      for (int qp = 0; qp < ir.GetNPoints(); qp++)
+      {
+         const IntegrationPoint &ip = ir.IntPoint(qp);
+         T->SetIntPoint(&ip);
+
+         Vector f(vdim);
+         f1_g.GetVectorValue(*T, ip, f);
+         for (int d = 0; d < vdim; d++)
+         {
+            int qpo = qp * vdim;
+            int eo = e * (ir.GetNPoints() * vdim);
+            f_test(d + qpo + eo) = f(d);
+         }
+      }
+   }
+
+   Vector diff(f_test);
+   diff -= y;
+   if (diff.Norml2() > 1e-10)
+   {
+      print_vector(diff);
+      print_vector(f_test);
+      print_vector(y);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_interpolate_linear_vector_3d);

examples/dfem_test_mass_scalar_2d.cpp

@@ -0,0 +1,111 @@
+#include "dfem/dfem_test_macro.hpp"
+#include "fem/bilininteg.hpp"
+#include <fstream>
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int dfem_test_mass_scalar_2d(std::string mesh_file,
+                             int refinements,
+                             int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   const int dim = mesh_serial.Dimension();
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   mesh_serial.Clear();
+
+   ParGridFunction *mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   ParGridFunction f1_g(&h1fes);
+
+   auto kernel_2d = [](const double &u,
+                       const tensor<double, 2, 2> &J,
+                       const double &w)
+   {
+      return mfem::tuple{u * w * det(J)};
+   };
+
+   auto kernel_3d = [](const double &u,
+                       const tensor<double, 3, 3> &J,
+                       const double &w)
+   {
+      return mfem::tuple{u * w * det(J)};
+   };
+
+   constexpr int Potential = 0;
+   constexpr int Coordinates = 1;
+
+   mfem::tuple input_operators = {Value<Potential>{}, Gradient<Coordinates>{}, Weight{}};
+   mfem::tuple output_operator = {Value<Potential>{}};
+
+   auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+   auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes}};
+
+   DifferentiableOperator dop(solutions, parameters, mesh);
+   if (dim == 2)
+   {
+      dop.AddDomainIntegrator(kernel_2d, input_operators, output_operator, ir);
+   }
+   else
+   {
+      dop.AddDomainIntegrator(kernel_3d, input_operators, output_operator, ir);
+   }
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      return 2.345 + x + x*y + 1.25 * x;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(*f1_g.GetTrueDofs()), y(h1fes.TrueVSize());
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+
+   ParBilinearForm a(&h1fes);
+   auto mass_integ = new MassIntegrator;
+   mass_integ->SetIntRule(&ir);
+   a.AddDomainIntegrator(mass_integ);
+   if (mesh.GetElement(0)->GetType() == Element::QUADRILATERAL ||
+       mesh.GetElement(0)->GetType() == Element::HEXAHEDRON)
+   {
+      a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   }
+   a.Assemble();
+   a.Finalize();
+
+   Vector y2(h1fes.TrueVSize());
+   a.Mult(x, y2);
+   y2.HostRead();
+
+   Vector diff(y2);
+   diff -= y;
+   // if (diff.Norml2() > 1e-12)
+   {
+      print_vector(diff);
+      print_vector(y2);
+      print_vector(y);
+      // return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(dfem_test_mass_scalar_2d);

examples/dfem_test_mass_scalar_3d.cpp

@@ -0,0 +1,147 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+#include "fem/bilininteg.hpp"
+#include "fem/fe/fe_base.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int dfem_test_mass_scalar_3d(std::string mesh_file,
+                             int refinements,
+                             int polynomial_order)
+{
+   constexpr int dim = 3;
+   Mesh mesh_serial = Mesh(mesh_file);
+   MFEM_ASSERT(mesh_serial.Dimension() == dim, "wrong mesh dimension");
+
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   ParGridFunction *mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   // IntegrationRules gll_rules(0, Quadrature1D::GaussLobatto);
+   // const IntegrationRule &ir = gll_rules.Get(h1fes.GetFE(0)->GetGeomType(),
+   //                                           2 * polynomial_order - 1);
+
+   auto dtq = h1fes.GetFE(0)->GetDofToQuad(ir, DofToQuad::TENSOR);
+   // printf("\n B: ");
+   // dtq.B.Print(out, dtq.B.Size());
+   // printf("\n G: ");
+   // dtq.G.Print(out, dtq.G.Size());
+   // printf("\n w: ");
+   // ir.GetWeights().Print(out, ir.GetWeights().Size());
+
+   // printf("#ndof per el = %d\n", h1fes.GetFE(0)->GetDof());
+   // printf("#nqp = %d\n", ir.GetNPoints());
+   // printf("#q1d = %d\n", (int)floor(pow(ir.GetNPoints(), 1.0/dim) + 0.5));
+
+   // printf("nodes: ");
+   // print_vector(*mesh_nodes);
+
+   ParGridFunction f1_g(&h1fes);
+
+   auto kernel = [](const double &u,
+                    const tensor<double, dim, dim> &J,
+                    const double &w)
+   {
+      return mfem::tuple{u * det(J) * w};
+   };
+
+   mfem::tuple argument_operators = {Value{"potential"}, Gradient{"coordinates"}, Weight{}};
+   mfem::tuple output_operator = {Value{"potential"}};
+
+   ElementOperator eop = {kernel, argument_operators, output_operator};
+   auto ops = mfem::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "potential"}};
+   auto parameters = std::array{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      const double z = coords(2);
+      return 2.345 + x + x*y + 1.25 * z*x;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+   // printf("\nf1_g: ");
+   // print_vector(f1_g);
+
+   auto R = h1fes.GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC);
+   // Vector f1_g_e(R->Height());
+   // R->Mult(f1_g, f1_g_e);
+   // printf("\nf1_g_e: ");
+   // print_vector(f1_g_e);
+   // auto r_out = std::ofstream("r_mat.mtx");
+   // R->PrintMatlab(r_out);
+   // r_out.close();
+
+   Vector x(*f1_g.GetTrueDofs()), y(h1fes.TrueVSize());
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+
+   ParBilinearForm a(&h1fes);
+   auto mass_integ = new MassIntegrator;
+   mass_integ->SetIntRule(&ir);
+   a.AddDomainIntegrator(mass_integ);
+   a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   a.Assemble();
+   a.Finalize();
+
+   Vector y2(h1fes.TrueVSize());
+   a.Mult(x, y2);
+   y2.HostRead();
+
+   Vector diff(y2);
+   diff -= y;
+   if (diff.Norml2() > 1e-15)
+   {
+      printf("y ");
+      print_vector(y);
+      printf("y2: ");
+      print_vector(y2);
+      printf("diff: ");
+      print_vector(diff);
+      return 1;
+   }
+
+   Vector y3(h1fes.TrueVSize());
+   auto dFdu = dop.GetDerivativeWrt<0>({&f1_g}, {mesh_nodes});
+
+   dFdu->Mult(x, y3);
+   diff = y2;
+   diff -= y;
+   if (diff.Norml2() > 1e-15)
+   {
+      printf("y2 ");
+      print_vector(y2);
+      printf("y3: ");
+      print_vector(y3);
+      printf("diff: ");
+      print_vector(diff);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(dfem_test_mass_scalar_3d);

examples/dfem_test_neo_hookean_elasticity_2d.cpp

@@ -0,0 +1,114 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_neo_hookean_elasticity_2d(
+   std::string mesh_file, int refinements, int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   MFEM_ASSERT(dim == 2, "This test is for 2D meshes only");
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction*>(mesh.GetNodes());
+   ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, dim);
+
+   out << "#dofs " << h1fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder());
+
+   out << "#qp: " << ir.GetNPoints() << "\n";
+
+   ParGridFunction u_g(&h1fes);
+
+   auto kernel = [] MFEM_HOST_DEVICE(const tensor<double, 2, 2>& J,
+                                     const double& w,
+                                     const tensor<double, 2, 2>& dudxi)
+   {
+      // Neo-Hookean parameters
+      const double lambda = 1.0;
+      const double mu = 0.5;
+
+      static constexpr auto I = mfem::internal::IsotropicIdentity<2>();
+      auto F = I + (dudxi * inv(J));
+      auto E = 0.5 * (transpose(F) * F - I);
+      auto invF = inv(F);
+
+      // 2D plane strain formulation
+      auto P = mu * (F - transpose(invF)) + lambda * log(det(F)) * transpose(invF);
+
+      return mfem::tuple{P * det(J) * w};
+   };
+
+   mfem::tuple argument_operators = {Gradient{"coordinates"}, Weight{},
+                                     Gradient{"displacement"}
+                                    };
+   mfem::tuple output_operator = {Gradient{"displacement"}};
+
+   ElementOperator eop = {kernel, argument_operators, output_operator};
+   auto ops = mfem::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "displacement"}};
+   auto parameters = std::array{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   auto displacement = [](const Vector& coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      u(0) = 0.1 * x * y;
+      u(1) = 0.1 * y * x;
+   };
+
+   VectorFunctionCoefficient disp_coeff(2, displacement);
+   u_g.ProjectCoefficient(disp_coeff);
+
+   Vector x(u_g), y(h1fes.TrueVSize());
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+   y.HostRead();
+
+   // Test linearization
+   auto dFdu = dop.GetDerivativeWrt<0>({&u_g}, {mesh_nodes});
+   dFdu->Mult(x, y);
+
+   // Finite difference Jacobian test
+   {
+      double eps = 1.0e-6;
+      Vector v(x), xpv(x), xmv(x), fxpv(x.Size()), fxmv(x.Size());
+      v *= eps;
+      xpv += v;
+      xmv -= v;
+      dop.Mult(xpv, fxpv);
+      dop.Mult(xmv, fxmv);
+      fxpv -= fxmv;
+      fxpv /= (2.0*eps);
+
+      fxpv -= y;
+      if (fxpv.Norml2() > eps)
+      {
+         out << "||dFdu_FD u^* - ex||_l2 = " << fxpv.Norml2() << "\n";
+         return 1;
+      }
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_neo_hookean_elasticity_2d);

examples/dfem_test_nonlinear_diffusion_3d.cpp

@@ -0,0 +1,169 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_nonlinear_diffusion(
+   std::string mesh_file, int refinements, int polynomial_order)
+{
+   constexpr int dim = 3;
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   mesh_serial.Clear();
+
+   out << "#el: " << mesh.GetNE() << "\n";
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction*>(mesh.GetNodes());
+   ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   out << "#dofs " << h1fes.GetTrueVSize() << "\n";
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   out << "#qp: " << ir.GetNPoints() << "\n";
+
+   ParGridFunction f1_g(&h1fes);
+
+   bool inactive_derivative = false;
+
+   auto kernel = [] MFEM_HOST_DEVICE(
+                    const tensor<double, dim, dim>& J,
+                    const double& w,
+                    const tensor<double, dim>& dudxi,
+                    const double& u)
+   {
+      auto invJ = inv(J);
+      return mfem::tuple{(u * u) * dudxi * invJ * transpose(invJ) * det(J) * w};
+   };
+
+   mfem::tuple argument_operators =
+   {
+      Gradient{"coordinates"},
+      Weight{},
+      Gradient{"potential"},
+      Value{"potential"}
+   };
+
+   mfem::tuple output_operator =
+   {
+      Gradient{"potential"}
+   };
+
+   ElementOperator eop = {kernel, argument_operators, output_operator};
+   auto ops = mfem::tuple{eop};
+
+   auto solutions = std::array
+   {
+      FieldDescriptor{&h1fes, "potential"}
+   };
+   auto parameters = std::array
+   {
+      FieldDescriptor{&mesh_fes, "coordinates"}
+   };
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   auto f1 = [](const Vector& coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      const double z = coords(2);
+      return 2.345 + 0.25 * x * x * y + y * y * x + z;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(f1_g), y(h1fes.TrueVSize());
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+   y.HostRead();
+
+   ParBilinearForm a(&h1fes);
+   GridFunctionCoefficient f1gc(&f1_g);
+   TransformedCoefficient tf_c(&f1gc, [](double f) { return f * f; });
+   a.AddDomainIntegrator(new DiffusionIntegrator(tf_c));
+   a.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   a.Assemble();
+   a.Finalize();
+
+   Vector y2(h1fes.TrueVSize()), diff(h1fes.TrueVSize());
+   a.Mult(x, y2);
+   y2.HostRead();
+   diff = y2;
+   diff -= y;
+   if (diff.Norml2() > 1e-10)
+   {
+      out << "||F(u) - ex||_l2 = " << diff.Norml2() << "\n";
+      print_vector(diff);
+      print_vector(y);
+      print_vector(y2);
+      return 1;
+   }
+
+   // Test linearization here as well
+   auto dFdu = dop.GetDerivativeWrt<0>({&f1_g}, {mesh_nodes});
+   dFdu->Mult(x, y);
+
+   // fd jacobian test
+   {
+      double eps = 1.0e-6;
+      Vector v(x), xpv(x), xmv(x), fxpv(x.Size()), fxmv(x.Size());
+      v *= eps;
+      xpv += v;
+      xmv -= v;
+      dop.Mult(xpv, fxpv);
+      dop.Mult(xmv, fxmv);
+      fxpv -= fxmv;
+      fxpv /= (2.0*eps);
+
+      fxpv -= y;
+      if (fxpv.Norml2() > eps)
+      {
+         out << "||dFdu_FD u^* - ex||_l2 = " << fxpv.Norml2() << "\n";
+         return 1;
+      }
+   }
+
+   // ParBilinearForm da(&h1fes);
+   // TransformedCoefficient dtf_c(&f1gc, [](double f) { return 2.0 * f; });
+   // da.AddDomainIntegrator(new DiffusionIntegrator(dtf_c));
+   // da.SetAssemblyLevel(AssemblyLevel::PARTIAL);
+   // da.Assemble();
+   // da.Finalize();
+
+   // if (dFdu->Height() != h1fes.GetTrueVSize())
+   // {
+   //    out << "dFdu unexpected height of " << dFdu->Height() << "\n";
+   //    return 1;
+   // }
+
+   // dFdu->Mult(x, y);
+   // print_vector(y);
+   // da.Mult(x, y2);
+   // print_vector(y2);
+   // y2 -= y;
+   // out << "||dFdu x - A x||_l2 = " << y2.Norml2() << "\n";
+   // if (y2.Norml2() > 1e-10)
+   // {
+   //    out << "||dFdu u^* - ex||_l2 = " << y2.Norml2() << "\n";
+   // }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_nonlinear_diffusion);

examples/dfem_test_nonlinear_elasticity_3d.cpp

@@ -0,0 +1,267 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+#include <fstream>
+
+using namespace mfem;
+using mfem::internal::tensor;
+using mfem::internal::dual;
+
+class FDJacobian : public Operator
+{
+public:
+   FDJacobian(const Operator &op, const Vector &x) :
+      Operator(op.Height()),
+      op(op),
+      x(x)
+   {
+      f.SetSize(Height());
+      xpev.SetSize(Height());
+      op.Mult(x, f);
+      xnorm = x.Norml2();
+   }
+
+   void Mult(const Vector &v, Vector &y) const override
+   {
+      x.HostRead();
+
+      // See [1] for choice of eps.
+      //
+      // [1] Woodward, C.S., Gardner, D.J. and Evans, K.J., 2015. On the use of
+      // finite difference matrix-vector products in Newton-Krylov solvers for
+      // implicit climate dynamics with spectral elements. Procedia Computer
+      // Science, 51, pp.2036-2045.
+      real_t eps = lambda * (lambda + xnorm / v.Norml2());
+
+      for (int i = 0; i < x.Size(); i++)
+      {
+         xpev(i) = x(i) + eps * v(i);
+      }
+
+      // y = f(x + eps * v)
+      op.Mult(xpev, y);
+
+      // y = (f(x + eps * v) - f(x)) / eps
+      for (int i = 0; i < x.Size(); i++)
+      {
+         y(i) = (y(i) - f(i)) / eps;
+      }
+   }
+
+   virtual MemoryClass GetMemoryClass() const override
+   {
+      return Device::GetDeviceMemoryClass();
+   }
+
+private:
+   const Operator &op;
+   Vector x, f;
+   mutable Vector xpev;
+   real_t lambda = 1.0e-6;
+   real_t xnorm;
+};
+
+template <typename elasticity_t>
+class ElasticityOperator : public Operator
+{
+   template <typename elasticity_du_t>
+   class ElasticityJacobianOperator : public Operator
+   {
+   public:
+      ElasticityJacobianOperator(const ElasticityOperator *elasticity,
+                                 std::shared_ptr<elasticity_du_t> dRdu) :
+         Operator(elasticity->Height()),
+         elasticity(elasticity),
+         dRdu(dRdu),
+         x_ess(dRdu->Height())
+      {
+      }
+
+      void Mult(const Vector &x, Vector &y) const override
+      {
+         x_ess = x;
+         x_ess.SetSubVector(elasticity->ess_tdofs, 0.0);
+
+         dRdu->Mult(x_ess, y);
+
+         for (int i = 0; i < elasticity->ess_tdofs.Size(); i++)
+         {
+            y[elasticity->ess_tdofs[i]] = x[elasticity->ess_tdofs[i]];
+         }
+      }
+
+      const ElasticityOperator *elasticity = nullptr;
+      std::shared_ptr<elasticity_du_t> dRdu;
+      mutable Vector x_ess;
+   };
+
+public:
+   ElasticityOperator(ParFiniteElementSpace &fes, elasticity_t &elasticity,
+                      Array<int> &ess_tdofs) :
+      Operator(fes.GetTrueVSize()),
+      fes(fes),
+      elasticity(elasticity),
+      ess_tdofs(ess_tdofs) {}
+
+   void Mult(const Vector &x, Vector &r) const override
+   {
+      elasticity.Mult(x, r);
+      r.SetSubVector(ess_tdofs, 0.0);
+   }
+
+   Operator &GetGradient(const Vector &x) const override
+   {
+      ParGridFunction u(const_cast<ParFiniteElementSpace *>
+                        (*std::get_if<const ParFiniteElementSpace *>
+                         (&elasticity.solutions[0].data)));
+
+      u.SetFromTrueDofs(x);
+      auto dRdu = elasticity.template GetDerivativeWrt<0>({&u}, {mesh_nodes});
+
+      jacobian.reset(
+         new ElasticityJacobianOperator<
+         typename std::remove_pointer<decltype(dRdu.get())>::type> (this, dRdu));
+
+      // jacobian.reset(new FDJacobian(*this, x));
+
+      return *jacobian;
+   }
+
+   void SetParameters(ParGridFunction &mesh_nodes)
+   {
+      elasticity.SetParameters({&mesh_nodes});
+      this->mesh_nodes = &mesh_nodes;
+   }
+
+   ParFiniteElementSpace &fes;
+   elasticity_t &elasticity;
+   Array<int> ess_tdofs;
+   mutable ParGridFunction *mesh_nodes = nullptr;
+   mutable std::shared_ptr<Operator> jacobian;
+};
+
+int test_nonlinear_elasticity_3d(std::string mesh_file,
+                                 int refinements,
+                                 int polynomial_order)
+{
+   constexpr int dim = 3;
+   constexpr int vdim = dim;
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   Array<int> ess_tdof_list, ess_bdr(mesh.bdr_attributes.Max());
+   ess_bdr = 0;
+   ess_bdr[0] = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list);
+
+   const IntegrationRule& ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(
+                      0)->GetDim() - 1);
+
+   out << "#qp: " << ir.GetNPoints() << "\n";
+   out << "#dof: " << h1fes.GetNDofs() << "\n";
+
+   ParGridFunction u(&h1fes);
+
+   auto elasticity_kernel = [] MFEM_HOST_DEVICE
+                            (const tensor<dual<real_t, real_t>, dim, dim> &dudxi,
+                             const tensor<real_t, dim, dim> &J,
+                             const real_t &w)
+   {
+      // shear modulus
+      real_t D1{0.1e6};
+      // bulk modulus
+      real_t C1{1.0e6};
+      constexpr auto I = mfem::internal::IsotropicIdentity<dim>();
+      auto invJ = inv(J);
+      auto dudx = dudxi * invJ;
+      auto F = det(I + dudx);
+      auto p = -2.0 * D1 * F * (F - 1);
+      auto devB = dev(dudx + transpose(dudx) + dot(dudx, transpose(dudx)));
+      auto sigma = -(p / F) * I + 2.0 * (C1 / pow(F, 5.0 / 3.0)) * devB;
+
+      return mfem::tuple{sigma * det(J) * w * transpose(invJ)};
+   };
+
+   mfem::tuple argument_operators{Gradient{"displacement"}, Gradient{"coordinates"}, Weight{}};
+   mfem::tuple output_operator{Gradient{"displacement"}};
+
+   // B^T D(B0*dudxi, B1*J, B2*w)
+   ElementOperator op(elasticity_kernel, argument_operators, output_operator, ir);
+
+   std::array solutions{FieldDescriptor{&h1fes, "displacement"}};
+   std::array parameters{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop(solutions, parameters, mfem::tuple{op}, mesh,
+                              AutoDiff::NativeDualNumber{});
+
+   ElasticityOperator elasticity(h1fes, dop, ess_tdof_list);
+
+   VectorArrayCoefficient f(dim);
+   for (int i = 0; i < dim-1; i++)
+   {
+      f.Set(i, new ConstantCoefficient(0.0));
+   }
+   {
+      Vector pull_force(mesh.bdr_attributes.Max());
+      pull_force = 0.0;
+      pull_force(1) = -1.0e-2;
+      f.Set(dim-1, new PWConstCoefficient(pull_force));
+   }
+
+   ParLinearForm b(&h1fes);
+   b.AddBoundaryIntegrator(new VectorBoundaryLFIntegrator(f));
+   b.UseFastAssembly(true);
+   b.Assemble();
+   auto B = b.ParallelAssemble();
+
+   Vector X = u.GetTrueVector();
+
+   CGSolver cg(MPI_COMM_WORLD);
+   cg.SetRelTol(1e-8);
+   cg.SetMaxIter(1000);
+   cg.SetPrintLevel(IterativeSolver::PrintLevel().Summary());
+
+   NewtonSolver newton(MPI_COMM_WORLD);
+   newton.SetSolver(cg);
+   newton.SetOperator(elasticity);
+   newton.SetRelTol(1e-6);
+   newton.SetMaxIter(100);
+   // newton.SetAdaptiveLinRtol();
+   newton.SetPrintLevel(IterativeSolver::PrintLevel().Iterations());
+
+   elasticity.SetParameters(*mesh_nodes);
+
+   // Vector zero;
+   newton.Mult(*B, X);
+
+   u.SetFromTrueDofs(X);
+
+   ParaViewDataCollection paraview_dc("dfem", &mesh);
+   paraview_dc.SetPrefixPath("ParaView");
+   paraview_dc.SetLevelsOfDetail(polynomial_order);
+   paraview_dc.SetDataFormat(VTKFormat::BINARY);
+   paraview_dc.SetHighOrderOutput(true);
+   paraview_dc.SetCycle(0);
+   paraview_dc.SetTime(0.0);
+   paraview_dc.RegisterField("displacement", &u);
+   paraview_dc.Save();
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_nonlinear_elasticity_3d);

examples/dfem_test_objective.cpp

@@ -0,0 +1,106 @@
+#include "dfem/dfem_test_macro.hpp"
+#include "examples/dfem/dfem_util.hpp"
+#include "fem/bilininteg.hpp"
+#include "fem/pbilinearform.hpp"
+#include "fem/plinearform.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_objective(
+   std::string mesh_file, int refinements, int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(),
+                   2 * h1fes.GetFE(0)->GetOrder() + h1fes.GetFE(0)->GetDim() - 1);
+
+   ParGridFunction f1_g(&h1fes);
+
+   auto kernel = [] MFEM_HOST_DEVICE (
+                    const real_t &u,
+                    const tensor<real_t, 2> &dudxi,
+                    const tensor<real_t, 2, 2> &J,
+                    const real_t &w)
+   {
+      return mfem::tuple{(u * u * norm(dudxi * inv(J))) * w * det(J)};
+   };
+
+   constexpr int Potential = 0;
+   constexpr int Coordinates = 1;
+
+   auto input_operators = mfem::tuple{Value<Potential>{}, Gradient<Potential>{}, Gradient<Coordinates>{}, Weight{}};
+   auto output_operator = mfem::tuple{One<Potential>{}};
+
+   auto solutions = std::vector{FieldDescriptor{Potential, &h1fes}};
+   auto parameters = std::vector{FieldDescriptor{Coordinates, &mesh_fes}};
+
+   DifferentiableOperator dop(solutions, parameters, mesh);
+   auto derivatives = std::integer_sequence<size_t, Potential> {};
+   dop.AddDomainIntegrator(kernel, input_operators, output_operator, ir,
+                           derivatives);
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      return x + y;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(*f1_g.GetTrueDofs());
+
+   Vector y(1);
+
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+
+   out << "Objective value ∫ u^2 dx (u = x+y) on reference element:\n";
+   print_vector(y);
+
+   auto dfdp = dop.GetDerivative(Potential, {&f1_g}, {mesh_nodes});
+   Vector dfdpv(1);
+   x = 1.0;
+   dfdp->Mult(x, dfdpv);
+
+   out << "Derivative of the objective wrt u:\n";
+   print_vector(dfdpv);
+
+   Vector dfdp_vec;
+   dfdp->AssembleVector(dfdp_vec);
+   out << "dfdp:\n";
+   print_vector(dfdp_vec);
+
+   {
+      x = *f1_g.GetTrueDofs();
+      FDJacobian fd_jac(dop, x);
+      x = 1.0;
+      fd_jac.Mult(x, y);
+      out << "fdjvp\n";
+      print_vector(y);
+      out << "FD Jacobian:\n";
+      fd_jac.PrintMatlab(out);
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_objective);

examples/dfem_test_objective_vjp.cpp

@@ -0,0 +1,107 @@
+#include "dfem/dfem_refactor.hpp"
+
+#include <fstream>
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+static constexpr int dim = 2;
+
+struct ObjectiveQFunction
+{
+    MFEM_HOST_DEVICE inline
+    auto operator()(const real_t& u, const tensor<real_t, dim, dim>& dXdxi, const double& w) const
+    {
+        auto dV = det(dXdxi)*w;
+        return mfem::tuple{u*dV};
+    }
+};
+
+int main()
+{
+    Mpi::Init();
+    out << std::setprecision(8);
+
+    int polynomial_order = 1;
+    int ir_order = 2;
+
+    Mesh mesh_serial = Mesh::MakeCartesian2D(1, 1, Element::QUADRILATERAL, false, 1.0, 1.0);
+    mesh_serial.EnsureNodes();
+    auto mesh = ParMesh(MPI_COMM_WORLD, mesh_serial);
+
+    auto mesh_nodes = static_cast<ParGridFunction*>(mesh.GetNodes());
+    ParFiniteElementSpace& mesh_fes = *mesh_nodes->ParFESpace();
+
+    out << "#el: " << mesh.GetNE() << "\n";
+
+    H1_FECollection fec(polynomial_order, dim);
+    ParFiniteElementSpace fes(&mesh, &fec);
+    const IntegrationRule &ir =
+      IntRules.Get(fes.GetFE(0)->GetGeomType(),
+                   2 * ir_order + fes.GetFE(0)->GetOrder());
+
+    ParGridFunction u(&fes);
+    //u = 0.0;
+    FunctionCoefficient fc([](const Vector& X) { return 2.0*X(0) - X(1); });
+    u.ProjectCoefficient(fc);
+    out << "u = ";
+    u.Print();
+
+    constexpr int Potential = 0;
+    constexpr int Coordinates = 1;
+
+    std::vector<FieldDescriptor> solutions{FieldDescriptor{Potential, &fes}};
+    std::vector<FieldDescriptor> parameters{FieldDescriptor{Coordinates, &mesh_fes}};
+    DifferentiableOperator op(solutions, parameters, mesh);
+    op.DisableTensorProductStructure();
+
+    // set up differentiable operator
+    ObjectiveQFunction qf;
+    mfem::tuple inputs{Value<Potential>{}, Gradient<Coordinates>{}, Weight{}};
+    // mfem::tuple outputs{One<Potential>{}};
+    mfem::tuple outputs{Value<Potential>{}};
+    Array<int> solid_domain_attr(mesh.attributes.Max());
+    solid_domain_attr[0] = 1;
+    auto derivatives = std::integer_sequence<size_t, Potential>{};
+    op.AddDomainIntegrator(qf, inputs, outputs, ir, solid_domain_attr, derivatives);
+    op.SetParameters({mesh_nodes});
+
+    Vector z(1);
+    op.Mult(u, z);
+    // z should be singleton, debug needed
+    out << "z = ";
+    z.Print();
+
+    // z_bar should be a singelton, need to debug
+    Vector z_bar(u.Size());
+    z_bar = 0.0;
+    z_bar(0) = 1.0;
+    out << "z_bar = ";
+    z_bar.Print();
+    Vector u_bar(u.Size());
+    auto jac = op.GetDerivative(Potential, {&u}, {mesh_nodes});
+    // jac->MultTranspose(z_bar, u_bar);
+    // out << "u_bar = ";
+    // u_bar.Print();
+
+    for (int i = 0; i < z_bar.Size(); i++) {
+        z_bar = 0.0;
+        z_bar(i) = 1.0;
+        jac->MultTranspose(z_bar, u_bar);
+        out << "jac col " << i << " = ";
+        u_bar.Print();
+    }
+
+    out << std::endl;
+    
+    Vector u_dot(u.Size());
+    Vector z_dot(u.Size());
+    for (int i = 0; i < z_bar.Size(); i++) {
+        u_dot = 0.0;
+        u_dot(i) = 1.0;
+        jac->Mult(u_dot, z_dot);
+        out << "jac row " << i << " = ";
+        z_dot.Print();
+    }
+    return 0;
+}

examples/dfem_test_ordering.cpp

@@ -0,0 +1,82 @@
+#include "dfem/dfem.hpp"
+#include "dfem/dfem_test_macro.hpp"
+#include "fem/coefficient.hpp"
+#include "fem/pgridfunc.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+int test_ordering(std::string mesh_file,
+                  int refinements,
+                  int polynomial_order)
+{
+   constexpr int dim = 2;
+   constexpr int vdim = dim;
+
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(polynomial_order);
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   const IntegrationRule &ir =
+      IntRules.Get(mesh_fes.GetFE(0)->GetGeomType(),
+                   2 * mesh_fes.FEColl()->GetOrder() - 1);
+
+   for (int q = 0; q < ir.GetNPoints(); q++)
+   {
+      out << "(" << ir.IntPoint(q).x << ", " << ir.IntPoint(q).y << ")\n";
+   }
+
+   ParGridFunction u(&mesh_fes);
+   auto f = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      u(0) = x*x*y + 1.0;
+      u(1) = y*y*x*x + 2.0;
+   };
+
+   VectorFunctionCoefficient uc(dim, f);
+   u.ProjectCoefficient(uc);
+
+   auto kernel = [](const tensor<double, dim> &xi,
+                    const tensor<double, vdim, dim> &J,
+                    const tensor<double, dim> &u,
+                    const tensor<double, vdim, dim> &dudxi)
+   {
+      out << "xi: " << xi << "\n";
+      out << "J: " << J << "\n";
+      out << "u: " << u << "\n";
+      out << "dudxi: " << dudxi << "\n\n";
+      return mfem::tuple{J};
+   };
+
+   mfem::tuple argument_operators{Value{"coordinates"}, Gradient{"coordinates"}, Value{"potential"}, Gradient{"potential"}};
+   mfem::tuple output_operator{Gradient{"potential"}};
+
+   ElementOperator op{kernel, argument_operators, output_operator};
+
+   std::array solutions{FieldDescriptor{&mesh_fes, "potential"}};
+   std::array parameters{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop{solutions, parameters, mfem::tuple{op}, mesh, ir};
+
+   Vector y(u);
+
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(u, y);
+
+   print_vector(y);
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_ordering);

examples/dfem_test_vector_diffusion.cpp

@@ -0,0 +1,117 @@
+#include "dfem/dfem_test_macro.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+
+template <int dim = 2>
+class VectorDiffusionQFunction
+{
+public:
+   VectorDiffusionQFunction() = default;
+
+   MFEM_HOST_DEVICE inline
+   auto operator() (const tensor<real_t, dim, dim>& dudxi,
+                    const tensor<real_t, dim, dim>& J,
+                    const real_t& w) const
+   {
+      auto invJ = inv(J);
+      return mfem::tuple{dudxi * invJ * det(J) * w * transpose(invJ)};
+   }
+};
+
+int test_vector_diffusion(std::string mesh_file,
+                          int refinements,
+                          int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   const int vdim = dim;
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, vdim);
+
+   Array<int> ess_bdr(mesh.bdr_attributes.Max());
+   Array<int> ess_tdof;
+   ess_bdr = 1;
+   h1fes.GetEssentialTrueDofs(ess_bdr, ess_tdof);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder() + 1);
+
+   ParGridFunction u(&h1fes);
+
+   auto f1 = [](const Vector& coords, Vector &u)
+   {
+      const real_t x = coords(0);
+      const real_t y = coords(1);
+      u(0) = 2.345 + 0.25 * x * x * y + y * y * x;
+      u(1) = 2.345 - 0.25 * x * y * y + y * x * x;
+   };
+
+   VectorFunctionCoefficient u_c(dim, f1);
+   u.ProjectCoefficient(u_c);
+
+   constexpr int Potential = 0;
+   constexpr int Coordinates = 1;
+
+   std::vector solutions{FieldDescriptor{Potential, &h1fes}};
+   std::vector parameters{FieldDescriptor{Coordinates, &mesh_fes}};
+   DifferentiableOperator dop{solutions, parameters, mesh};
+
+   VectorDiffusionQFunction vector_diffusion_kernel;
+   mfem::tuple input_operators{Gradient<Potential>{}, Gradient<Coordinates>{}, Weight{}};
+   mfem::tuple output_operator{Gradient<Potential>{}};
+   auto derivatives = std::integer_sequence<size_t, Potential> {};
+   dop.AddDomainIntegrator(vector_diffusion_kernel, input_operators,
+                           output_operator, ir, derivatives);
+
+   Vector x(u), y1(h1fes.GetTrueVSize()), y2(h1fes.GetTrueVSize());
+
+   ParBilinearForm A_form(&h1fes);
+   auto A_integ = new VectorDiffusionIntegrator(vdim);
+   A_integ->SetIntegrationRule(ir);
+   A_form.AddDomainIntegrator(A_integ);
+   A_form.Assemble();
+   A_form.Finalize();
+
+   HypreParMatrix *A_mfem = A_form.ParallelAssemble();
+   A_mfem->PrintMatlab(out);
+   out << "\n";
+
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y1);
+   y1.HostRead();
+
+   HypreParMatrix A_dfem;
+   dop.GetDerivative(Potential, {&u}, {mesh_nodes})->Assemble(A_dfem);
+   A_dfem.PrintMatlab(out);
+
+   A_form.Mult(x, y2);
+   y2.HostRead();
+
+   Vector diff(y2);
+   diff -= y1;
+   if (diff.Norml2() > 1e-10)
+   {
+      out << "||F(u) - ex||_l2 = " << diff.Norml2() << "\n";
+      print_vector(diff);
+      print_vector(y1);
+      print_vector(y2);
+      return 1;
+   }
+
+   return 0;
+}
+
+DFEM_TEST_MAIN(test_vector_diffusion);

examples/enzyme_interface_smoketest.cpp

@@ -0,0 +1,122 @@
+#include <tuple>
+#include <type_traits>
+#include <iostream>
+#include <enzyme/enzyme>
+
+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);
+}
+
+template <typename ... Ts>
+constexpr auto decay_types(std::tuple<Ts...> const &)
+-> std::tuple<std::remove_cv_t<std::remove_reference_t<Ts>>...>;
+
+template <typename T>
+using decay_tuple = decltype(decay_types(std::declval<T>()));
+
+template <class F> struct FunctionSignature;
+
+template <typename output_t, typename... input_ts>
+struct FunctionSignature<output_t(input_ts...)>
+{
+   using return_t = output_t;
+   using parameter_ts = std::tuple<input_ts...>;
+};
+
+template <class T> struct create_function_signature;
+
+template <typename output_t, typename T, typename... input_ts>
+struct create_function_signature<output_t (T::*)(input_ts...) const>
+{
+   using type = FunctionSignature<output_t(input_ts...)>;
+};
+
+template <typename arg_ts, std::size_t... Is>
+auto create_enzyme_args(arg_ts &args,
+                        arg_ts &shadow_args,
+                        std::index_sequence<Is...>)
+{
+   // (std::cout << ... << std::get<Is>(shadow_args));
+   return std::tuple<enzyme::Duplicated<decltype(std::get<Is>(args))>...>
+   {
+      { std::get<Is>(args), std::get<Is>(shadow_args) }...
+   };
+}
+
+template <typename kernel_t, typename arg_ts>
+auto fwddiff_apply_enzyme(kernel_t kernel, arg_ts &&args, arg_ts &&shadow_args)
+{
+   auto arg_indices =
+      std::make_index_sequence<std::tuple_size_v<std::remove_reference_t<arg_ts>>> {};
+
+   auto enzyme_args = create_enzyme_args(args, shadow_args, arg_indices);
+
+   // using kf_return_t = typename create_function_signature<
+   //                     decltype(&kernel_t::operator())>::type::return_t;
+
+   std::cout << "\n";
+   std::cout << "args is " << get_type_name<decltype(args)>() << "\n\n";
+   std::cout << "enzyme_args type is " << get_type_name<decltype(enzyme_args)>() <<
+             "\n\n";
+   // std::cout << "return type is " << get_type_name<decltype(kf_return_t{})>() <<
+   //           "\n\n";
+
+   std::cout << "args " << std::get<0>(args) << "\n";
+   std::cout << "shadow args " << std::get<0>(shadow_args) << "\n";
+
+   return std::apply([&](auto &&...args)
+   {
+      // std::cout << enzyme::autodiff<enzyme::Forward>(+kernel, args...) << "\n";
+      return enzyme::get<0>
+             (enzyme::autodiff<enzyme::Forward>(+kernel, args...));
+   },
+   enzyme_args);
+}
+
+int main()
+{
+
+   auto func = [](const double &x, double &y)
+   {
+      std::cout << "func( x = " << x << " )\n";
+      return x*x;
+   };
+
+   using kf_param_ts = typename create_function_signature<
+                       decltype(&decltype(func)::operator())>::type::parameter_ts;
+   using kf_output_t = typename create_function_signature<
+                       decltype(&decltype(func)::operator())>::type::return_t;
+   auto kernel_args = decay_tuple<kf_param_ts> {};
+   auto kernel_shadow_args = decay_tuple<kf_param_ts> {};
+
+   std::get<0>(kernel_args) = 3;
+   std::get<0>(kernel_shadow_args) = 1;
+
+   auto dx = fwddiff_apply_enzyme(func, kernel_args, kernel_shadow_args);
+
+   std::cout << "dfdx = " << dx << "\n";
+
+   return 0;
+}

examples/test_dfem_dual.cpp

@@ -0,0 +1,305 @@
+#include "dfem/dfem_refactor.hpp"
+#include "linalg/hypre.hpp"
+
+using namespace mfem;
+using mfem::internal::tensor;
+using mfem::internal::dual;
+
+int test_diffusion_integrator(std::string mesh_file,
+                              int refinements,
+                              int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder());
+
+   ParGridFunction f1_g(&h1fes);
+   ParGridFunction rho_g(&h1fes);
+
+   auto rho_f = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      return x + y;
+   };
+
+   FunctionCoefficient rho_c(rho_f);
+   rho_g.ProjectCoefficient(rho_c);
+
+   auto kernel = [](const tensor<dual<double, double>, 2> &grad_u,
+                    const dual<double, double> &rho,
+                    const tensor<double, 2, 2> &J,
+                    const double &w)
+   {
+      auto invJ = inv(J);
+      return std::tuple{rho*rho * grad_u * invJ * transpose(invJ) * det(J) * w};
+   };
+
+   std::tuple argument_operators = {Gradient{"potential"}, Value{"density"}, Gradient{"coordinates"}, Weight{"integration_weights"}};
+   std::tuple output_operator = {Gradient{"potential"}};
+
+   ElementOperator eop = {kernel, argument_operators, output_operator};
+   auto ops = std::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "potential"}};
+   auto parameters = std::array
+   {
+      FieldDescriptor{&h1fes, "density"},
+      FieldDescriptor{&mesh_fes, "coordinates"}
+   };
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   auto f1 = [](const Vector &coords)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      return 2.345 + 0.25 * x*x*y + y*y*x;
+   };
+
+   FunctionCoefficient f1_c(f1);
+   f1_g.ProjectCoefficient(f1_c);
+
+   Vector x(f1_g), y(h1fes.TrueVSize());
+   dop.SetParameters({&rho_g, mesh_nodes});
+   dop.Mult(x, y);
+
+   ParBilinearForm a(&h1fes);
+   TransformedCoefficient rho_c2(&rho_c, [](double c) {return c*c;});
+   a.AddDomainIntegrator(new DiffusionIntegrator(rho_c2));
+   a.Assemble();
+   a.Finalize();
+
+   Vector y2(h1fes.TrueVSize());
+   a.Mult(x, y2);
+   y2 -= y;
+   if (y2.Norml2() > 1e-10)
+   {
+      out << "||F(u) - ex||_l2 = " << y2.Norml2() << "\n";
+      return 1;
+   }
+
+   // Test linearization here as well
+   auto dFdu = dop.GetDerivativeWrt<0>({&f1_g}, {&rho_g, mesh_nodes});
+
+   // HypreParMatrix A;
+   // dFdu->Assemble(A);
+
+   if (dFdu->Height() != h1fes.GetTrueVSize())
+   {
+      out << "dFdu unexpected height of " << dFdu->Height() << "\n";
+      return 1;
+   }
+
+   dFdu->Mult(x, y);
+   a.Mult(x, y2);
+   y2 -= y;
+   if (y2.Norml2() > 1e-10)
+   {
+      out << "||dFdu u^* - ex||_l2 = " << y2.Norml2() << "\n";
+      return 1;
+   }
+
+   // fd jacobian test
+   {
+      double eps = 1.0e-6;
+      Vector v(x), xpv(x), xmv(x), fxpv(x.Size()), fxmv(x.Size());
+      v *= eps;
+      xpv += v;
+      xmv -= v;
+      dop.Mult(xpv, fxpv);
+      dop.Mult(xmv, fxmv);
+      fxpv -= fxmv;
+      fxpv /= (2.0*eps);
+
+      fxpv -= y;
+      if (fxpv.Norml2() > eps)
+      {
+         out << "||dFdu_FD u^* - ex||_l2 = " << fxpv.Norml2() << "\n";
+         return 1;
+      }
+   }
+
+   f1_g.ProjectCoefficient(f1_c);
+   rho_g.ProjectCoefficient(rho_c);
+   auto dFdrho = dop.GetDerivativeWrt<1>({&f1_g}, {&rho_g, mesh_nodes});
+   if (dFdrho->Height() != h1fes.GetTrueVSize())
+   {
+      out << "dFdrho unexpected height of " << dFdrho->Height() << "\n";
+      return 1;
+   }
+
+   dFdrho->Mult(rho_g, y);
+
+   // fd test
+   {
+      double eps = 1.0e-6;
+      Vector v(rho_g), rhopv(rho_g), rhomv(rho_g), frhopv(x.Size()), frhomv(x.Size());
+      v *= eps;
+      rhopv += v;
+      rhomv -= v;
+      dop.SetParameters({&rhopv, mesh_nodes});
+      dop.Mult(x, frhopv);
+      dop.SetParameters({&rhomv, mesh_nodes});
+      dop.Mult(x, frhomv);
+      frhopv -= frhomv;
+      frhopv /= (2.0*eps);
+
+      frhopv -= y;
+      if (frhopv.Norml2() > eps)
+      {
+         out << "||dFdu_FD u^* - ex||_l2 = " << frhopv.Norml2() << "\n";
+         return 1;
+      }
+   }
+
+   return 0;
+}
+
+int test_qoi(std::string mesh_file,
+             int refinements,
+             int polynomial_order)
+{
+   Mesh mesh_serial = Mesh(mesh_file);
+   for (int i = 0; i < refinements; i++)
+   {
+      mesh_serial.UniformRefinement();
+   }
+   ParMesh mesh(MPI_COMM_WORLD, mesh_serial);
+
+   mesh.SetCurvature(1);
+   const int dim = mesh.Dimension();
+   mesh_serial.Clear();
+
+   ParGridFunction* mesh_nodes = static_cast<ParGridFunction *>(mesh.GetNodes());
+   ParFiniteElementSpace &mesh_fes = *mesh_nodes->ParFESpace();
+
+   H1_FECollection h1fec(polynomial_order, dim);
+   ParFiniteElementSpace h1fes(&mesh, &h1fec, dim);
+
+   const IntegrationRule &ir =
+      IntRules.Get(h1fes.GetFE(0)->GetGeomType(), 2 * h1fec.GetOrder());
+
+   ParGridFunction rho_g(&h1fes);
+
+   auto rho_f = [](const Vector &coords, Vector &u)
+   {
+      const double x = coords(0);
+      const double y = coords(1);
+      u(0) = x + y;
+      u(1) = x + y;
+   };
+
+   VectorFunctionCoefficient rho_c(dim, rho_f);
+   rho_g.ProjectCoefficient(rho_c);
+
+   auto kernel = [](const tensor<dual<double, double>, 2> &rho,
+                    const tensor<dual<double, double>, 2, 2> &drhodxi,
+                    const tensor<double, 2, 2> &J,
+                    const double &w)
+   {
+      const double eps = 1.2345;
+      const auto drhodx = drhodxi * inv(J);
+      return std::tuple{(0.5 * eps * dot(rho, rho) + ddot(drhodx, drhodx)) * det(J) * w};
+   };
+
+   std::tuple argument_operators = {Value{"density"}, Gradient{"density"}, Gradient{"coordinates"}, Weight{"integration_weights"}};
+   std::tuple output_operator = {One{"density"}};
+
+   ElementOperator eop = {kernel, argument_operators, output_operator};
+   auto ops = std::tuple{eop};
+
+   auto solutions = std::array{FieldDescriptor{&h1fes, "density"}};
+   auto parameters = std::array{FieldDescriptor{&mesh_fes, "coordinates"}};
+
+   DifferentiableOperator dop(solutions, parameters, ops, mesh, ir);
+
+   Vector x(rho_g), y(1);
+   dop.SetParameters({mesh_nodes});
+   dop.Mult(x, y);
+
+   // print_vector(y);
+
+   auto dFdrho = dop.GetDerivativeWrt<0>({&rho_g}, {mesh_nodes});
+   // Vector dFdrho_vec;
+   // dFdrho->Assemble(dFdrho_vec);
+
+   // print_vector(dFdrho_vec);
+
+   // fd jacobian test
+   {
+      double eps = 1.0e-8;
+      Vector v(x), fxpv(1), fxmv(1), dfdx(x.Size());
+      for (int i = 0; i < x.Size(); i++)
+      {
+         v(i) += eps;
+         dop.Mult(v, fxpv);
+         v(i) -= 2.0 * eps;
+         dop.Mult(v, fxmv);
+         fxpv -= fxmv;
+         fxpv /= (2.0*eps);
+         dfdx(i) = fxpv(0);
+      }
+
+      // print_vector(dfdx);
+      dfdx -= dFdrho_vec;
+      if (dfdx.Norml2() > 1e-6)
+      {
+         out << "||dFdu_FD u^* - ex||_l2 = " << dfdx.Norml2() << "\n";
+         return 1;
+      }
+   }
+
+   return 0;
+}
+
+int main(int argc, char *argv[])
+{
+   Mpi::Init();
+
+   std::cout << std::setprecision(9);
+
+   const char *mesh_file = "../data/star.mesh";
+   int polynomial_order = 1;
+   int ir_order = 2;
+   int refinements = 0;
+
+   OptionsParser args(argc, argv);
+   args.AddOption(&mesh_file, "-m", "--mesh", "Mesh file to use.");
+   args.AddOption(&polynomial_order, "-o", "--order", "");
+   args.AddOption(&refinements, "-r", "--r", "");
+   args.AddOption(&ir_order, "-iro", "--iro", "");
+   args.ParseCheck();
+
+   out << std::setprecision(12);
+
+   int ret;
+
+   ret = test_diffusion_integrator(mesh_file,
+                                   refinements,
+                                   polynomial_order);
+   out << "test_diffusion_integrator";
+   ret ? out << " FAILURE\n" : out << " OK\n";
+
+   ret = test_qoi(mesh_file, refinements, polynomial_order);
+   out << "test_qoi";
+   ret ? out << " FAILURE\n" : out << " OK\n";
+
+   return 0;
+}

general/enzyme.hpp

@@ -21,16 +21,16 @@
  * documentation (https://enzyme.mit.edu) for more information.
  */
 
-extern int enzyme_dup;
-extern int enzyme_dupnoneed;
-extern int enzyme_out;
-extern int enzyme_const;
+// extern int enzyme_dup;
+// extern int enzyme_dupnoneed;
+// extern int enzyme_out;
+// extern int enzyme_const;
 
-template <typename return_type, typename... Args>
-return_type __enzyme_autodiff(Args...);
+// template <typename return_type, typename... Args>
+// return_type __enzyme_autodiff(Args...);
 
-template <typename return_type, typename... Args>
-return_type __enzyme_fwddiff(Args...);
+// template <typename return_type, typename... Args>
+// return_type __enzyme_fwddiff(Args...);
 
 #define MFEM_ENZYME_INACTIVENOFREE   __attribute__((enzyme_inactive, enzyme_nofree))
 #define MFEM_ENZYME_INACTIVE   __attribute__((enzyme_inactive))

general/mem_manager.hpp

@@ -871,6 +871,8 @@ public:
    inline static void* __enzyme_allocation_like2[4] = {(void*)static_cast<void*(*)(void*, size_t, MemoryType, MemoryType, unsigned, unsigned&)>(MemoryManager::New_),
                                                        (void*)1, (void*)"-1,2,4", (void*)MemoryManager::Delete_
                                                       };
+   __attribute__((used))
+   inline static void* __enzyme_function_like[2] = {(void*)MemoryManager::Delete_, (void*)"free"};
 #endif
 };
 

linalg/dtensor.hpp

@@ -9,6 +9,129 @@
 // terms of the BSD-3 license. We welcome feedback and contributions, see file
 // CONTRIBUTING.md for details.
 
+// #ifndef MFEM_DTENSOR
+// #define MFEM_DTENSOR
+
+// #include <numeric>
+
+// #include "../config/config.hpp"
+
+// namespace mfem
+// {
+
+// /// ///////////////////////////////////////////////////////////////////////////
+// template <int N, typename... Args>
+// MFEM_HOST_DEVICE inline int ColMajor(const int (&dims)[N], Args... args)
+// {
+//    int offset = 0, i = 0, unused;
+//    ((offset *= dims[N - ++i], offset += args, std::ignore = unused) = ...);
+//    return offset;
+// }
+
+// /// ///////////////////////////////////////////////////////////////////////////
+// template <int N, typename... Args>
+// MFEM_HOST_DEVICE inline int RowMajor(const int (&dims)[N], Args... args)
+// {
+//    int offset = 0, i = 0;
+//    return ((..., (offset *= dims[i++], offset += args, 0)), offset);
+// }
+
+// /// ///////////////////////////////////////////////////////////////////////////
+// template <int N, typename T = real_t, bool Column = true>
+// class DeviceTensor
+// {
+// protected:
+//    int dims[N], size;
+//    T *data;
+
+// public:
+//    /// Default constructor
+//    // DeviceTensor() = delete;
+//    MFEM_HOST_DEVICE
+//    DeviceTensor() {}
+
+//    /// Constructor to initialize a tensor from the Scalar array data_
+//    template <typename... Args, typename = std::enable_if_t<(sizeof...(Args) == N)>>
+//    MFEM_HOST_DEVICE DeviceTensor(T *data, Args... args):
+//       dims{args...},
+//       size(std::accumulate(dims, dims + N, 1, std::multiplies<int> {})),
+//         data(size > 0 ? data : nullptr)
+//    { }
+
+//    /// Copy constructor (default)
+//    // DeviceTensor(const DeviceTensor &) = default;
+
+//    /// Copy assignment (default)
+//    // DeviceTensor &operator=(const DeviceTensor &) = default;
+
+//    /// Conversion to `Scalar *`.
+//    MFEM_HOST_DEVICE inline operator T *() const { return data; }
+
+//    /// Computes the offset of the tensor element at the given multi-indices
+//    template <typename... Args, typename = std::enable_if_t<(sizeof...(Args) == N)>>
+//    MFEM_HOST_DEVICE inline int Offset(Args... args) const
+//    {
+//       constexpr auto offset = Column ? ColMajor<N, Args...> : RowMajor<N, Args...>;
+//       return offset(dims, args...);
+//    }
+
+//    /// Const accessor for the data
+//    template <typename... Args, typename = std::enable_if_t<(sizeof...(Args) == N)>>
+//    MFEM_HOST_DEVICE inline T &operator()(Args... args) const
+//    {
+//       return data[Offset(args...)];
+//    }
+
+//    /// Subscript operator where the tensor is viewed as a 1D array.
+//    MFEM_HOST_DEVICE inline T &operator[](int i) const { return data[i]; }
+
+//    /// Returns the size of the tensor
+//    MFEM_HOST_DEVICE inline int Size() const { return size; }
+
+//    MFEM_HOST_DEVICE inline std::array<int, N> GetShape() const
+//    {
+//       std::array<int, N> s;
+//       for (int i = 0; i < N; i++)
+//       {
+//          s[i] = dims[i];
+//       }
+//       return s;
+//    }
+// };
+
+// /** @brief Wrap a pointer as a DeviceTensor with automatically deduced template
+//     parameters */
+// template <typename T, typename... Dims>
+// MFEM_HOST_DEVICE inline DeviceTensor<sizeof...(Dims), T> Reshape(T *ptr,
+//                                                                  Dims... dims)
+// {
+//    return DeviceTensor<sizeof...(Dims), T>(ptr, dims...);
+// }
+
+// template <typename T, typename... Dims>
+// MFEM_HOST_DEVICE inline DeviceTensor<sizeof...(Dims), T, false> RowReshape(
+//    T *ptr,
+//    Dims... dims)
+// {
+//    return DeviceTensor<sizeof...(Dims), T, false>(ptr, dims...);
+// }
+
+// using DeviceArray = DeviceTensor<1, int>;
+// using ConstDeviceArray = DeviceTensor<1, const int>;
+
+// using DeviceVector = DeviceTensor<1, real_t>;
+// using ConstDeviceVector = DeviceTensor<1, const real_t>;
+
+// using DeviceMatrix = DeviceTensor<2, real_t>;
+// using ConstDeviceMatrix = DeviceTensor<2, const real_t>;
+
+// using DeviceCube = DeviceTensor<3, real_t>;
+// using ConstDeviceCube = DeviceTensor<3, const real_t>;
+
+// } // namespace mfem
+
+// #endif // MFEM_DTENSOR
+
 #ifndef MFEM_DTENSOR
 #define MFEM_DTENSOR
 
@@ -87,7 +210,9 @@ protected:
 
 public:
    /// Default constructor
-   DeviceTensor() = delete;
+   // DeviceTensor() = delete;
+   MFEM_HOST_DEVICE
+   DeviceTensor() {}
 
    /// Constructor to initialize a tensor from the Scalar array data_
    template <typename... Args> MFEM_HOST_DEVICE
@@ -122,6 +247,16 @@ public:
    {
       return data[i];
    }
+
+   MFEM_HOST_DEVICE inline std::array<int, Dim> GetShape() const
+   {
+      std::array<int, Dim> s;
+      for (int i = 0; i < Dim; i++)
+      {
+         s[i] = sizes[i];
+      }
+      return s;
+   }
 };
 
 

linalg/tensor.hpp

@@ -19,6 +19,8 @@
 #define MFEM_INTERNAL_TENSOR_HPP
 
 #include "dual.hpp"
+#include "general/backends.hpp"
+#include <limits>
 #include <type_traits> // for std::false_type
 
 namespace mfem
@@ -436,6 +438,23 @@ tensor<decltype(f(n1, n2, n3, n4)), n1, n2, n3, n4>
    return A;
 }
 
+// needs to be generalized
+template <typename T, int m, int n> MFEM_HOST_DEVICE
+tensor<T, n> get_col(tensor<T, m, n> A, int j)
+{
+   tensor<T, n> c{};
+   c(0) = A[0][j];
+   c(1) = A[1][j];
+   return c;
+}
+
+/// @overload
+template <typename T> MFEM_HOST_DEVICE
+tensor<T, 1> get_col(tensor<T, 1, 1> A, int j)
+{
+   return tensor<T, 1> {A[0][0]};
+}
+
 /**
  * @brief return the sum of two tensors
  * @tparam S the underlying type of the lefthand argument
@@ -697,6 +716,20 @@ auto outer(S A, T B) -> decltype(A * B)
    return A * B;
 }
 
+template <typename T, int n, int m> MFEM_HOST_DEVICE
+tensor<T, n + m> flatten(tensor<T, n, m> A)
+{
+   tensor<T, n + m> B{};
+   for (int i = 0; i < n; i++)
+   {
+      for (int j = 0; j < m; j++)
+      {
+         B(i + j * m) = A(i, j);
+      }
+   }
+   return B;
+}
+
 /**
  * @overload
  * @note this overload implements the case where the left argument is a scalar, and the right argument is a tensor
@@ -1051,6 +1084,18 @@ decltype(S {} * T{})
    return AB;
 }
 
+template <typename T, int m> MFEM_HOST_DEVICE
+auto dot(const tensor<T, m>& A, const tensor<T, m>& B) ->
+decltype(T {})
+{
+   decltype(T{}) AB{};
+   for (int i = 0; i < m; i++)
+   {
+      AB += A[i] * B[i];
+   }
+   return AB;
+}
+
 template <typename S, typename T, int m, int... n> MFEM_HOST_DEVICE
 auto dot(const tensor<S, m>& A, const tensor<T, m, n...>& B) ->
 tensor<decltype(S {} * T{}), n...>
@@ -1321,6 +1366,12 @@ tensor<T, n, m> transpose(const tensor<T, m, n>& A)
  * @param[in] A The matrix to obtain the determinant of
  */
 template <typename T> MFEM_HOST_DEVICE
+T det(const tensor<T, 1, 1>& A)
+{
+   return A[0][0];
+}
+/// @overload
+template <typename T> MFEM_HOST_DEVICE
 T det(const tensor<T, 2, 2>& A)
 {
    return A[0][0] * A[1][1] - A[0][1] * A[1][0];
@@ -1335,6 +1386,145 @@ T det(const tensor<T, 3, 3>& A)
           A[2][0];
 }
 
+template <typename T> MFEM_HOST_DEVICE
+std::tuple<tensor<T, 1>, tensor<T, 1, 1>> eig(tensor<T, 1, 1> &A)
+{
+   return {tensor<T, 1>{A[0][0]}, tensor<T, 1, 1>{{1.0}}};
+}
+
+template <typename T> MFEM_HOST_DEVICE
+std::tuple<tensor<T, 2>, tensor<T, 2, 2>> eig(tensor<T, 2, 2> &A)
+{
+   tensor<T, 2> e;
+   tensor<T, 2, 2> v;
+
+   double d0 = A(0, 0);
+   double d2 = A(0, 1);
+   double d3 = A(1, 1);
+   double c, s;
+
+   if (d2 == 0.0)
+   {
+      c = 1.0;
+      s = 0.0;
+   }
+   else
+   {
+      double t;
+      const double zeta = (d3 - d0) / (2.0 * d2);
+      const double azeta = fabs(zeta);
+      if (azeta < std::sqrt(1.0/std::numeric_limits<T>::epsilon()))
+      {
+         t = copysign(1./(azeta + std::sqrt(1. + zeta*zeta)), zeta);
+      }
+      else
+      {
+         t = copysign(0.5/azeta, zeta);
+      }
+      c = std::sqrt(1./(1. + t*t));
+      s = c*t;
+      t *= d2;
+      d0 -= t;
+      d3 += t;
+   }
+
+   if (d0 <= d3)
+   {
+      e(0) = d0;
+      e(1) = d3;
+      v(0, 0) =  c;
+      v(1, 0) = -s;
+      v(0, 1) =  s;
+      v(1, 1) =  c;
+   }
+   else
+   {
+      e(0) = d3;
+      e(1) = d0;
+      v(0, 0) =  s;
+      v(1, 0) =  c;
+      v(0, 1) =  c;
+      v(1, 1) = -s;
+   }
+
+   return {e, v};
+}
+
+template <typename T> MFEM_HOST_DEVICE
+void GetScalingFactor(const T &d_max, T &mult)
+{
+   int d_exp;
+   if (d_max > 0.)
+   {
+      mult = frexp(d_max, &d_exp);
+      if (d_exp == std::numeric_limits<T>::max_exponent)
+      {
+         mult *= std::numeric_limits<T>::radix;
+      }
+      mult = d_max/mult;
+   }
+   else
+   {
+      mult = 1.;
+   }
+}
+
+template <typename T> MFEM_HOST_DEVICE
+T calcsv(const tensor<T, 1, 1> A, const int i)
+{
+   return A[0][0];
+}
+
+/**
+ * @brief Compute the i-th singular value of a 2x2 matrix A
+ */
+template <typename T> MFEM_HOST_DEVICE
+T calcsv(const tensor<T, 2, 2> A, const int i)
+{
+   double mult;
+   double d0, d1, d2, d3;
+   d0 = A(0, 0);
+   d1 = A(1, 0);
+   d2 = A(0, 1);
+   d3 = A(1, 1);
+
+   double d_max = fabs(d0);
+   if (d_max < fabs(d1)) { d_max = fabs(d1); }
+   if (d_max < fabs(d2)) { d_max = fabs(d2); }
+   if (d_max < fabs(d3)) { d_max = fabs(d3); }
+
+   GetScalingFactor(d_max, mult);
+
+   d0 /= mult;
+   d1 /= mult;
+   d2 /= mult;
+   d3 /= mult;
+
+   double t = 0.5*((d0+d2)*(d0-d2)+(d1-d3)*(d1+d3));
+   double s = d0*d2 + d1*d3;
+   s = std::sqrt(0.5*(d0*d0 + d1*d1 + d2*d2 + d3*d3) + std::sqrt(t*t + s*s));
+
+   if (s == 0.0)
+   {
+      return 0.0;
+   }
+   t = fabs(d0*d3 - d1*d2) / s;
+   if (t > s)
+   {
+      if (i == 0)
+      {
+         return t*mult;
+      }
+      return s*mult;
+   }
+   if (i == 0)
+   {
+      return s*mult;
+   }
+   return t*mult;
+}
+
+
 /**
  * @brief Return whether a square rank 2 tensor is symmetric
  *
@@ -1474,13 +1664,20 @@ tensor<T, n> linear_solve(tensor<T, n, n> A, const tensor<T, n> b)
 /**
  * @brief Inverts a matrix
  * @param[in] A The matrix to invert
- * @note Uses a shortcut for inverting a 2-by-2 matrix
+ * @note Uses a shortcut for inverting a 1x1, 2x2 and 3x3 matrix
  */
-inline MFEM_HOST_DEVICE tensor<real_t, 2, 2> inv(const tensor<real_t, 2, 2>& A)
+template <typename T>
+inline MFEM_HOST_DEVICE tensor<T, 1, 1> inv(const tensor<T, 1, 1>& A)
 {
-   real_t inv_detA(1.0 / det(A));
+   return tensor<T, 1, 1> {T{1.0} / A[0][0]};
+}
 
-   tensor<real_t, 2, 2> invA{};
+template <typename T>
+inline MFEM_HOST_DEVICE tensor<T, 2, 2> inv(const tensor<T, 2, 2>& A)
+{
+   T inv_detA(1.0 / det(A));
+
+   tensor<T, 2, 2> invA{};
 
    invA[0][0] = A[1][1] * inv_detA;
    invA[0][1] = -A[0][1] * inv_detA;
@@ -1494,11 +1691,12 @@ inline MFEM_HOST_DEVICE tensor<real_t, 2, 2> inv(const tensor<real_t, 2, 2>& A)
  * @overload
  * @note Uses a shortcut for inverting a 3-by-3 matrix
  */
-inline MFEM_HOST_DEVICE tensor<real_t, 3, 3> inv(const tensor<real_t, 3, 3>& A)
+template <typename T>
+inline MFEM_HOST_DEVICE tensor<T, 3, 3> inv(const tensor<T, 3, 3>& A)
 {
-   real_t inv_detA(1.0 / det(A));
+   T inv_detA(1.0 / det(A));
 
-   tensor<real_t, 3, 3> invA{};
+   tensor<T, 3, 3> invA{};
 
    invA[0][0] = (A[1][1] * A[2][2] - A[1][2] * A[2][1]) * inv_detA;
    invA[0][1] = (A[0][2] * A[2][1] - A[0][1] * A[2][2]) * inv_detA;
@@ -1520,7 +1718,7 @@ inline MFEM_HOST_DEVICE tensor<real_t, 3, 3> inv(const tensor<real_t, 3, 3>& A)
 template <typename T, int n> MFEM_HOST_DEVICE
 tensor<T, n, n> inv(const tensor<T, n, n>& A)
 {
-   auto abs  = [](real_t x) { return (x < 0) ? -x : x; };
+   auto abs  = [](T x) { return (x < 0) ? -x : x; };
    auto swap = [](tensor<T, n>& x, tensor<T, n>& y)
    {
       auto tmp = x;
@@ -1528,12 +1726,12 @@ tensor<T, n, n> inv(const tensor<T, n, n>& A)
       y        = tmp;
    };
 
-   tensor<real_t, n, n> B = Identity<n>();
+   tensor<T, n, n> B = Identity<n>();
 
    for (int i = 0; i < n; i++)
    {
       // Search for maximum in this column
-      real_t max_val = abs(A[i][i]);
+      T max_val = abs(A[i][i]);
 
       int max_row = i;
       for (int j = i + 1; j < n; j++)
@@ -1553,7 +1751,7 @@ tensor<T, n, n> inv(const tensor<T, n, n>& A)
       {
          if (A[j][i] != 0.0)
          {
-            real_t c = -A[j][i] / A[i][i];
+            T c = -A[j][i] / A[i][i];
             A[j] += c * A[i];
             B[j] += c * B[i];
             A[j][i] = 0;

sum_factorization.jl

@@ -0,0 +1,190 @@
+using LinearAlgebra
+using Tullio
+
+num_qp = 4
+num_trial_dof = 4;
+trial_vdim = 1
+trial_op_dim = [1, 2]
+test_op_dim = 2
+test_vdim = 1
+num_test_dof = num_trial_dof
+num_rho_dof = 4;
+space_dim = 2;
+
+Bu_mem = [0.622008468 0.166666667 0.166666667 0.0446581987 0.166666667 0.622008468 0.0446581987 0.166666667 0.0446581987 0.166666667 0.166666667 0.622008468 0.166666667 0.0446581987 0.622008468 0.166666667]
+
+Bdu_mem = [-0.788675135 -0.788675135 -0.211324865 -0.211324865 -0.788675135 -0.211324865 -0.788675135 -0.211324865 0.788675135 0.788675135 0.211324865 0.211324865 -0.211324865 -0.788675135 -0.211324865 -0.788675135 0.211324865 0.211324865 0.788675135 0.788675135 0.211324865 0.788675135 0.211324865 0.788675135 -0.211324865 -0.211324865 -0.788675135 -0.788675135 0.788675135 0.211324865 0.788675135 0.211324865]
+
+Bu = reshape(Bu_mem, (num_qp, trial_op_dim[1], num_trial_dof))
+Bdu = reshape(Bdu_mem, (num_qp, trial_op_dim[2], num_trial_dof))
+Bv = Bdu;
+
+u_e = reshape([2.345 2.345 3.595 2.345], (num_trial_dof, trial_vdim))
+rho_e = reshape([0.0 1.0 2.0 1.0], (num_trial_dof, trial_vdim))
+x_e = reshape([0.0 1.0 1.0 0.0 0.0 0.0 1.0 1.0], (num_trial_dof, space_dim))
+
+rho_qpref = [0.422650 1.000000 1.000000 1.577350]
+J_qpref = [1.000000 0.000000 0.000000 1.000000 1.000000 0.000000 0.000000 1.000000 1.000000 0.000000 0.000000 1.000000 1.000000 0.000000 0.000000 1.000000]
+w_qpref = [0.250000 0.250000 0.250000 0.250000]
+dudxi_qpref = [0.264156 0.264156 0.264156 0.985844 0.985844 0.264156 0.985844 0.985844]
+
+rho_qp = zeros(Float64, trial_vdim, num_qp)
+for v = 1:trial_vdim
+    for q = 1:num_qp
+        acc = 0.0
+        for d = 1:num_trial_dof
+            acc += Bu[q, v, d] * rho_e[d, v]
+        end
+        rho_qp[v, q] = acc
+    end
+end
+println(rho_qp)
+
+J_qp = zeros(Float64, space_dim, space_dim, num_qp)
+for v = 1:space_dim
+    for s = 1:space_dim
+        for q = 1:num_qp
+            acc = 0.0
+            for d = 1:num_trial_dof
+                acc += Bdu[q, s, d] * x_e[d, v]
+            end
+            J_qp[v, s, q] = acc
+        end
+    end
+end
+println(J_qp)
+
+dudxi_qp = zeros(Float64, trial_vdim, space_dim, num_qp)
+for v = 1:trial_vdim
+    for s = 1:space_dim
+        for q = 1:num_qp
+            acc = 0.0
+            for d = 1:num_trial_dof
+                acc += Bdu[q, s, d] * u_e[d, v]
+            end
+            dudxi_qp[v, s, q] = acc
+        end
+    end
+end
+println(dudxi_qp)
+
+w_qp = w_qpref
+
+# sum factorization
+
+u_e = reshape(Float64[2.345 2.345 2.345 3.595], (num_trial_dof, trial_vdim))
+rho_e = reshape(Float64[0 1 1 2], (num_trial_dof, trial_vdim))
+x_e = reshape(Float64[0 1 0 1 0 0 1 1], (num_trial_dof, space_dim))
+
+nq1d = 2
+nd1d = 2
+
+B = reshape([0.788675135 0.211324865 0.211324865 0.788675135], (nq1d, nd1d))
+G = reshape([-1 -1 1 1], (nq1d, nd1d))
+
+function interpolate_value(u_e)
+    u_e = reshape(u_e, (nd1d, nd1d))
+    S2 = zeros(Float64, nq1d, nq1d)
+    for v = 1:trial_vdim
+        @tullio S1[qx, dy] := u_e[dx, dy] * B[qx, dx]
+        @tullio S2[qx, qy] = B[qy, dy] * S1[qx, dy]
+    end
+    return S2
+end
+
+function interpolate_grad(u_e)
+    vdim = size(u_e, 2)
+    u_e = reshape(u_e, (nd1d, nd1d, vdim))
+    dq0 = zeros(Float64, nd1d, nq1d)
+    dq1 = zeros(Float64, nd1d, nq1d)
+    dudxi_qp = zeros(nq1d, nq1d, vdim, space_dim)
+
+    for vd = 1:vdim
+        for dy = 1:nd1d
+            for qx = 1:nq1d
+                u = 0.0
+                v = 0.0
+                for dx = 1:nd1d
+                    u += u_e[dx, dy, vd] * B[qx, dx]
+                    v += u_e[dx, dy, vd] * G[qx, dx]
+                end
+                dq0[dy, qx] = u
+                dq1[dy, qx] = v
+            end
+        end
+
+        for qy = 1:nq1d
+            for qx = 1:nq1d
+                du = [0.0, 0.0]
+                for dy = 1:nd1d
+                    du[1] += dq1[dy, qx] * B[qy, dy]
+                    du[2] += dq0[dy, qx] * G[qy, dy]
+                end
+
+                for s = 1:space_dim
+                    dudxi_qp[qx, qy, vd, s] = du[s]
+                end
+            end
+        end
+    end
+    return dudxi_qp
+end
+
+function integrate_grad(f_qp, r)
+    dq0 = zeros(Float64, nd1d, nq1d)
+    dq1 = zeros(Float64, nd1d, nq1d)
+    for qy = 1:nq1d
+        for dx = 1:nd1d
+            u = 0.0
+            v = 0.0
+            for qx = 1:nq1d
+                u += G[qx, dx] * f_qp[qx, qy, 1, 1]
+                v += B[qx, dx] * f_qp[qx, qy, 1, 2]
+            end
+            dq0[dx, qy] = u
+            dq1[dx, qy] = v
+        end
+    end
+
+    for dy = 1:nd1d
+        for dx = 1:nd1d
+            u = 0.0
+            v = 0.0
+            for qy = 1:nq1d
+                u += dq0[dx, qy] * B[qy, dy]
+                v += dq1[dx, qy] * G[qy, dy]
+            end
+            r[dx, dy] += u + v
+        end
+    end
+end
+
+function integrate_value(f_qp, r)
+end
+
+function kernel(rho, J, w, dudxi)
+    return rho * rho * ((inv(J) * transpose(inv(J))) * dudxi) * det(J) * w
+end
+
+rho_qp = reshape(interpolate_value(rho_e), (nq1d * nq1d,))
+dudxi_qp = reshape(interpolate_grad(u_e), (nq1d * nq1d, 2))
+J_qp = reshape(interpolate_grad(x_e), (nq1d * nq1d), 2, 2)
+f_qp = zeros(nq1d * nq1d, 1, space_dim)
+for q = 1:nq1d*nq1d
+    f_qp[q, 1, :] = kernel(rho_qp[q], J_qp[q, :, :], w_qp[q], dudxi_qp[q, :])
+end
+
+w_qp = reshape(w_qp, nq1d, nq1d)
+rho_qp = interpolate_value(rho_e)
+dudxi_qp = reshape(interpolate_grad(u_e), (nq1d, nq1d, 2))
+J_qp = interpolate_grad(x_e)
+f_qp = zeros(nq1d, nq1d, 1, space_dim)
+for qx = 1:nq1d
+    for qy = 1:nq1d
+        f_qp[qx, qy, 1, :] = kernel(rho_qp[qx, qy], J_qp[qx, qy, :, :], w_qp[qx, qy], dudxi_qp[qx, qy, :])
+    end
+end
+
+r = zeros(Float64, nd1d, nd1d)
+integrate_grad(f_qp, r)
+println(reshape(r, nd1d * nd1d))

sum_factorization_3d.jl

@@ -0,0 +1,285 @@
+using LinearAlgebra
+using DelimitedFiles
+
+Q = 4
+D = 3
+space_dim = 3;
+num_trial_dof = D^3;
+test_vdim = 1
+output_op_dim = 3;
+
+R_data = Int.(readdlm("/Users/andrej1/repos/mfem/build-debug/r_mat.mtx", ' ', Float64))
+R_N = maximum(R_data[:, 1:2])
+R = zeros(Float64, (R_N, R_N))
+for e = 1:size(R_data, 1)
+    rij = R_data[e, :, :]
+    R[rij[1], rij[2]] = 1.0
+end
+
+u_l = reshape(Float64[2.345 3.345 4.345 2.345 2.345 4.595 5.595 2.345 2.845 3.845 3.345 2.345 3.47 5.095 3.97 2.345 2.345 3.97 4.97 2.345 3.095 3.1575 4.47 3.6575 2.345 3.72 3.4075], (num_trial_dof, test_vdim))
+u_e = R * u_l
+x_l = transpose(reshape(Float64[0 0 0 1 0 0 1 1 0 0 1 0 0 0 1 1 0 1 1 1 1 0 1 1 0.5 0 0 1 0.5 0 0.5 1 0 0 0.5 0 0.5 0 1 1 0.5 1 0.5 1 1 0 0.5 1 0 0 0.5 1 0 0.5 1 1 0.5 0 1 0.5 0.5 0.5 0 0.5 0 0.5 1 0.5 0.5 0.5 1 0.5 0 0.5 0.5 0.5 0.5 1 0.5 0.5 0.5], (space_dim, num_trial_dof)))
+x_e = R * x_l
+
+w_qp = [0.00526143468632 0.00986393947438 0.00986393947438 0.00526143468632 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.00526143468632 0.00986393947438 0.00986393947438 0.00526143468632 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.0184925420071 0.0346691208692 0.0346691208692 0.0184925420071 0.0184925420071 0.0346691208692 0.0346691208692 0.0184925420071 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.0184925420071 0.0346691208692 0.0346691208692 0.0184925420071 0.0184925420071 0.0346691208692 0.0346691208692 0.0184925420071 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.00526143468632 0.00986393947438 0.00986393947438 0.00526143468632 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.00986393947438 0.0184925420071 0.0184925420071 0.00986393947438 0.00526143468632 0.00986393947438 0.00986393947438 0.00526143468632]
+B = reshape([0.80134602937 0.227784076791 -0.112196966794 -0.0597902822241 0.258444252854 0.884412890003 0.884412890003 0.258444252854 -0.0597902822241 -0.112196966794 0.227784076791 0.80134602937], (Q, D))
+G = reshape([-2.72227262319 -1.67996208717 -0.32003791283 0.722272623188 3.44454524638 1.35992417434 -1.35992417434 -3.44454524638 -0.722272623188 0.32003791283 1.67996208717 2.72227262319], (Q, D))
+
+function interpolate_value(f_e)
+    vdim = size(f_e, 2)
+    f_e = reshape(f_e, (D, D, D, vdim))
+    f_qp = zeros(vdim, Q, Q, Q)
+    s1 = zeros(Float64, D, D, Q)
+    s2 = zeros(Float64, D, Q, Q)
+
+    for vd = 1:vdim
+        for dz = 1:D
+            for dy = 1:D
+                for qx = 1:Q
+                    acc = 0.0
+                    for dx = 1:D
+                        acc += f_e[dx, dy, dz, vd] * B[qx, dx]
+                    end
+                    s1[dz, dy, qx] = acc
+                end
+            end
+        end
+
+        for dz = 1:D
+            for qx = 1:Q
+                for qy = 1:Q
+                    acc = 0.0
+                    for dy = 1:D
+                        acc += s1[dz, dy, qx] * B[qy, dy]
+                    end
+                    s2[dz, qy, qx] = acc
+                end
+            end
+        end
+
+        for qz = 1:Q
+            for qy = 1:Q
+                for qx = 1:Q
+                    acc = 0.0
+                    for dz = 1:D
+                        acc += s2[dz, qy, qx] * B[qz, dz]
+                    end
+                    f_qp[vd, qx, qy, qz] = acc
+                end
+            end
+        end
+    end
+
+    return f_qp
+end
+
+function interpolate_grad(f_e)
+    vdim = size(f_e, 2)
+    f_e = reshape(f_e, (D, D, D, vdim))
+    dudxi_qp = zeros(vdim, space_dim, Q, Q, Q)
+
+    s1 = zeros(Float64, D, D, Q)
+    s2 = zeros(Float64, D, D, Q)
+    s3 = zeros(Float64, D, Q, Q)
+    s4 = zeros(Float64, D, Q, Q)
+    s5 = zeros(Float64, D, Q, Q)
+    uvw = zeros(Float64, 3)
+
+    for vd = 1:vdim
+        for dz = 1:D
+            for dy = 1:D
+                for qx = 1:Q
+                    uvw .= 0.0
+                    for dx = 1:D
+                        f = f_e[dx, dy, dz, vd]
+                        uvw[1] += f * B[qx, dx]
+                        uvw[2] += f * G[qx, dx]
+                    end
+                    s1[dz, dy, qx] = uvw[1]
+                    s2[dz, dy, qx] = uvw[2]
+                end
+            end
+        end
+
+        for dz = 1:D
+            for qy = 1:Q
+                for qx = 1:Q
+                    uvw .= 0.0
+                    for dy = 1:D
+                        uvw[1] += s2[dz, dy, qx] * B[qy, dy]
+                        uvw[2] += s1[dz, dy, qx] * G[qy, dy]
+                        uvw[3] += s1[dz, dy, qx] * B[qy, dy]
+                    end
+                    s3[dz, qy, qx] = uvw[1]
+                    s4[dz, qy, qx] = uvw[2]
+                    s5[dz, qy, qx] = uvw[3]
+                end
+            end
+        end
+
+        for qz = 1:Q
+            for qy = 1:Q
+                for qx = 1:Q
+                    uvw .= 0.0
+                    for dz = 1:D
+                        uvw[1] += s3[dz, qy, qx] * B[qz, dz]
+                        uvw[2] += s4[dz, qy, qx] * B[qz, dz]
+                        uvw[3] += s5[dz, qy, qx] * G[qz, dz]
+                    end
+                    dudxi_qp[vd, 1, qx, qy, qz] = uvw[1]
+                    dudxi_qp[vd, 2, qx, qy, qz] = uvw[2]
+                    dudxi_qp[vd, 3, qx, qy, qz] = uvw[3]
+                end
+            end
+        end
+    end
+
+    return dudxi_qp
+end
+
+function integrate_value(f_qp)
+    r = zeros(Float64, D, D, D, test_vdim)
+    f_qp = reshape(f_qp, (test_vdim, output_op_dim, Q, Q, Q))
+    s1 = zeros(Float64, Q, Q, D)
+    s2 = zeros(Float64, Q, D, D)
+
+    for vd = 1:test_vdim
+        for qy = 1:Q
+            for dx = 1:D
+                for qz = 1:Q
+                    acc = 0.0
+                    for qx = 1:Q
+                        acc += f_qp[test_vdim, 1, qx, qy, qz] * B[qx, dx]
+                    end
+                    s1[qz, qy, dx] = acc
+                end
+            end
+        end
+
+        # for i = 1:Q
+        #     for j = 1:Q
+        #         for k = 1:D
+        #             print(s1[i,j,k], " ")
+        #         end
+        #     end
+        # end
+        # println()
+
+        for dy = 1:D
+            for dx = 1:D
+                for qz = 1:Q
+                    acc = 0.0
+                    for qy = 1:Q
+                        acc += s1[qz, qy, dx] * B[qy, dy]
+                    end
+                    s2[qz, dy, dx] = acc
+                end
+            end
+        end
+
+        for dy = 1:D
+            for dx = 1:D
+                for dz = 1:D
+                    acc = 0.0
+                    for qz = 1:Q
+                        acc += s2[qz, dy, dx] * B[qz, dz]
+                    end
+                    r[dx, dy, dz, vd] += acc
+                end
+
+            end
+        end
+    end
+
+    return r
+end
+
+function integrate_grad(f_qp)
+    r = zeros(Float64, D, D, D, test_vdim)
+    f_qp = reshape(f_qp, (test_vdim, output_op_dim, Q, Q, Q))
+    s0 = zeros(Float64, Q, Q, D)
+    s1 = zeros(Float64, Q, Q, D)
+    s2 = zeros(Float64, Q, Q, D)
+    s3 = zeros(Float64, Q, D, D)
+    s4 = zeros(Float64, Q, D, D)
+    s5 = zeros(Float64, Q, D, D)
+
+    for vd = 1:test_vdim
+        for qz = 1:Q
+            for qy = 1:Q
+                for dx = 1:D
+                    uvw = zeros(Float64, 3)
+                    for qx = 1:Q
+                        uvw[1] += f_qp[vd, 1, qx, qy, qz] * G[qx, dx]
+                        uvw[2] += f_qp[vd, 2, qx, qy, qz] * B[qx, dx]
+                        uvw[3] += f_qp[vd, 3, qx, qy, qz] * B[qx, dx]
+                    end
+                    s0[qz, qy, dx] = uvw[1]
+                    s1[qz, qy, dx] = uvw[2]
+                    s2[qz, qy, dx] = uvw[3]
+                end
+            end
+        end
+
+        for qz = 1:Q
+            for dy = 1:D
+                for dx = 1:D
+                    uvw = zeros(Float64, 3)
+                    for qy = 1:Q
+                        uvw[1] += s0[qz, qy, dx] * B[qy, dy]
+                        uvw[2] += s1[qz, qy, dx] * G[qy, dy]
+                        uvw[3] += s2[qz, qy, dx] * B[qy, dy]
+                    end
+                    s3[qz, dy, dx] = uvw[1]
+                    s4[qz, dy, dx] = uvw[2]
+                    s5[qz, dy, dx] = uvw[3]
+                end
+            end
+        end
+
+        for dz = 1:D
+            for dy = 1:D
+                for dx = 1:D
+                    uvw = zeros(Float64, 3)
+                    for qz = 1:Q
+                        uvw[1] += s3[qz, dy, dx] * B[qz, dz]
+                        uvw[2] += s4[qz, dy, dx] * B[qz, dz]
+                        uvw[3] += s5[qz, dy, dx] * G[qz, dz]
+                    end
+                    r[dx, dy, dz, vd] += sum(uvw)
+                end
+            end
+        end
+
+        for i = 1:D
+            for j = 1:D
+                for k = 1:D
+                    print(f_qp[1, 1, i,j,k], " ")
+                end
+            end
+        end
+        println()
+    end
+
+    return r
+end
+
+function kernel(dudxi, J, w)
+    transpose(inv(J)) * inv(J) * dudxi * det(J) * w
+end
+
+dudxi_qp = reshape(interpolate_grad(u_e), (space_dim, Q * Q * Q))
+J_qp = reshape(interpolate_grad(x_e), (space_dim, space_dim, Q * Q * Q))
+f_qp = zeros(test_vdim, output_op_dim, Q * Q * Q)
+
+for q = 1:Q*Q*Q
+    r_qp = kernel(dudxi_qp[:, q], J_qp[:, :, q], w_qp[q])
+    for od = 1:output_op_dim
+        f_qp[1, od, q] += r_qp[od]
+    end
+end
+
+r = integrate_grad(f_qp)
+# println(transpose(R) * reshape(r, (D * D * D)))
+
+println(reshape(r, D * D * D))

transfermaps.jl

@@ -0,0 +1,145 @@
+using Enzyme, ForwardDiff, LinearAlgebra
+
+num_qp = 4;
+num_trial_dof = 4;
+trial_vdim = 1
+trial_op_dim = [1, 2]
+num_test_dof = num_trial_dof
+test_op_dim = 2
+test_vdim = 1
+
+Bu_u_mem = [0.622008468 0.166666667 0.166666667 0.0446581987 0.166666667 0.622008468 0.0446581987 0.166666667 0.0446581987 0.166666667 0.166666667 0.622008468 0.166666667 0.0446581987 0.622008468 0.166666667]
+Bu_du_mem = [-0.788675135 -0.788675135 -0.211324865 -0.211324865 -0.788675135 -0.211324865 -0.788675135 -0.211324865 0.788675135 0.788675135 0.211324865 0.211324865 -0.211324865 -0.788675135 -0.211324865 -0.788675135 0.211324865 0.211324865 0.788675135 0.788675135 0.211324865 0.788675135 0.211324865 0.788675135 -0.211324865 -0.211324865 -0.788675135 -0.788675135 0.788675135 0.211324865 0.788675135 0.211324865]
+Bu_u = reshape(Bu_u_mem, (num_qp, trial_op_dim[1], num_trial_dof))
+Bu_du = reshape(Bu_du_mem, (num_qp, trial_op_dim[2], num_trial_dof))
+Bu = [Bu_u, Bu_du]
+Bv = Bu_du;
+
+det(A) = A[1, 1] * A[2, 2] - A[1, 2] * A[2, 1]
+
+# kernel(u, J, w) = u * det(J) * w
+
+# u = [1.0, 1.0];
+# du = [0.0, 0.0];
+J = [1 0; 0 1];
+w = 0.25;
+
+### ParamtricFunction test
+
+# ParamtricFunction as quadrature data
+pf_size_on_qp = 4
+pf = zeros(num_qp * pf_size_on_qp)
+
+residual_size_on_qp = zeros(length(pf))
+
+function kernel(J, w)
+    return J * w
+end
+
+function Trho(u, e, q, size_on_qp, num_qp)
+    b = q * size_on_qp
+    c = q * size_on_qp + ((e - 1) * num_qp * size_on_qp)
+    return u[b-size_on_qp+1:c]
+end
+
+function TrhoT(u, uq, e, q, size_on_qp, num_qp)
+    b = q * size_on_qp
+    c = q * size_on_qp + ((e - 1) * num_qp * size_on_qp)
+    return u[b-size_on_qp+1:c] = uq
+end
+
+for e = 1:1
+    for q = 1:num_qp
+        pfq = Trho(pf, e, q, pf_size_on_qp, num_qp)
+        pfq[:] = kernel(J, w)
+        TrhoT(pf, pfq, e, q, pf_size_on_qp, num_qp)
+    end
+end
+
+# println(pf)
+
+nqp = 4
+data = reshape([1 1 1 1 0 0 0 0 0 0 0 0 1 1 1 1], (16,))
+m = 2
+n = 2
+arg = zeros(n, m)
+for q = 1:nqp
+    for i = 1:m
+        for j = 1:n
+            arg[j, i] = data[(i * m) + j]
+        end
+    end
+    println(arg)
+end
+
+
+# function kernel(u, dudxi, J, w)
+#     invJ = J
+#     dudx = dudxi * invJ
+#     return u[1, 1] * dudx * det(J) * w * transpose(invJ)
+# end
+
+# u = ones((1, 1))
+# du = [zeros((1, 1)), zeros((2, 2))]
+# dudxi = zeros((2, 2))
+# J = [1 0; 0 1];
+# w = 0.25;
+
+# wrap(x) = kernel(x, J, w)
+
+# D = zeros(test_vdim, test_op_dim, trial_vdim, sum(trial_op_dim), num_qp)
+# for q = 1:num_qp
+#     for j = 1:trial_vdim
+#         m_offset = 0
+#         local trial_op_dim = size(Bu[s])[2]
+#         for s = 1:length(Bu)
+#             for m = 1:trial_op_dim
+#                 du[s][j, m] = 1.0
+#                 df = autodiff(Forward, kernel, Duplicated(u, du[1]), Duplicated(dudxi, du[2]), Const(J), Const(w))[1]
+#                 du[s][j, m] = 0.0
+#                 for i = 1:test_vdim
+#                     for k = 1:test_op_dim
+#                         D[i, k, j, m+m_offset, q] = df[i, k]
+#                     end
+#                 end
+#             end
+#             m_offset += trial_op_dim
+#         end
+#     end
+# end
+
+# Ae = zeros(num_test_dof, test_vdim, num_trial_dof, trial_vdim)
+
+# for J = 1:num_trial_dof
+#     for j = 1:trial_vdim
+#         fhat = zeros(test_vdim, test_op_dim, num_qp)
+#         m_offset = 0
+#         for s = 1:length(Bu)
+#             trial_op_dim = size(Bu[s])[2]
+#             # precompute fhat for trial dof J column
+#             for qp = 1:num_qp
+#                 for i = 1:test_vdim
+#                     for k = 1:test_op_dim
+#                         for m = 1:trial_op_dim
+#                             fhat[i, k, qp] += D[i, k, j, m+m_offset, qp] * Bu[s][qp, m, J]
+#                         end
+#                     end
+#                 end
+#             end
+#             m_offset += trial_op_dim
+#         end
+
+#         # this imitates what 'map_quadrature_data_to_fields' does
+#         for I = 1:num_test_dof
+#             for i = 1:test_vdim
+#                 for qp = 1:num_qp
+#                     for k = 1:test_op_dim
+#                         Ae[I, i, J, j] += fhat[i, k, qp] * Bv[qp, k, I]
+#                     end
+#                 end
+#             end
+#         end
+#     end
+# end
+
+# display(reshape(Ae, (num_test_dof * test_vdim, num_trial_dof * trial_vdim)))