Enable field vjp in dfem

src/smith/physics/dfem_weak_form.hpp

@@ -150,6 +150,8 @@ class DfemWeakForm : public WeakForm {
   {
     // sum field operator doesn't work with sum factorization
     v_dot_weak_form_residual_.DisableTensorProductStructure();
+    // use TVECTORs where we can in dFEM
+    weak_form_.SetMultLevel(mfem::future::DifferentiableOperator::MultLevel::TVECTOR);
     residual_vector_.UseDevice(true);
   }
 
@@ -205,16 +207,45 @@ class DfemWeakForm : public WeakForm {
 
   /// @overload
   std::unique_ptr<mfem::HypreParMatrix> jacobian(double /*time*/, double dt, ConstFieldPtr /*shape_disp*/,
-                                                 const std::vector<ConstFieldPtr>& /*fields*/,
-                                                 const std::vector<double>& /*jacobian_weights*/,
+                                                 const std::vector<ConstFieldPtr>& fields,
+                                                 const std::vector<double>& jacobian_weights,
                                                  const std::vector<ConstQuadratureFieldPtr>& /*quad_fields*/ = {},
-                                                 int /*block_row*/ = 0) const override
+                                                 int block_row = 0) const override
   {
-    SLIC_ERROR_ROOT("DfemWeakForm does not support matrix assembly");
+    SLIC_ERROR_ROOT_IF(block_row != 0, "Invalid block row and column requested in fieldJacobian for DfemWeakForm");
 
     dt_ = dt;
 
-    return std::make_unique<mfem::HypreParMatrix>();
+    std::unique_ptr<mfem::HypreParMatrix> J;
+
+    auto addToJ = [&J](double factor, std::unique_ptr<mfem::HypreParMatrix> jac_contrib) {
+      if (J) {
+        SLIC_ERROR_IF(J->N() != jac_contrib->N(),
+                      "Multiple nonzero jacobian weights are being used on inconsistently sized input arguments.");
+        SLIC_ERROR_IF(J->M() != jac_contrib->M(),
+                      "Multiple nonzero jacobian weights are being used on inconsistently sized input arguments.");
+        J->Add(factor, *jac_contrib);
+      } else {
+        J.reset(jac_contrib.release());
+        if (factor != 1.0) (*J) *= factor;
+      }
+    };
+
+    mfem::Vector empty_residual_l(output_mfem_space_.GetVSize());
+    std::vector<mfem::Vector*> test_par_gf({&empty_residual_l});
+    std::vector<mfem::Vector*> field_par_gf = getLVectors(fields);
+
+    for (size_t input_col = 0; input_col < jacobian_weights.size(); ++input_col) {
+      if (jacobian_weights[input_col] != 0.0) {
+        auto deriv_op = weak_form_.GetDerivative(input_col, test_par_gf, field_par_gf);
+        auto jac_contrib = std::make_unique<mfem::HypreParMatrix>();
+        auto jac_contrib_ptr = jac_contrib.get();
+        deriv_op->Assemble(jac_contrib_ptr);
+        addToJ(jacobian_weights[input_col], std::move(jac_contrib));
+      }
+    }
+
+    return J;
   }
 
   /// @overload
@@ -247,39 +278,26 @@ class DfemWeakForm : public WeakForm {
   }
 
   /// @overload
-  void vjp(double /*time*/, double dt, ConstFieldPtr /*shape_disp*/, const std::vector<ConstFieldPtr>& /*fields*/,
-           const std::vector<ConstQuadratureFieldPtr>& /*quad_fields*/, const std::vector<ConstFieldPtr>& /*v_fields*/,
-           DualFieldPtr /*vjp_shape_disp_sensitivity*/, const std::vector<DualFieldPtr>& /*vjp_sensitivities*/,
+  void vjp(double /*time*/, double dt, ConstFieldPtr /*shape_disp*/, const std::vector<ConstFieldPtr>& fields,
+           const std::vector<ConstQuadratureFieldPtr>& /*quad_fields*/, const std::vector<ConstFieldPtr>& v_fields,
+           DualFieldPtr /*vjp_shape_disp_sensitivity*/, const std::vector<DualFieldPtr>& vjp_sensitivities,
            const std::vector<QuadratureFieldPtr>& /*vjp_quad_field_sensitivities*/) const override
   {
-    SLIC_ERROR_ROOT("DfemWeakForm does not support vjp calculations");
-
-    // SLIC_ERROR_IF(vjp_sensitivities.size() != fields.size(),
-    //               "Invalid number of field sensitivities relative to the number of fields");
-    // SLIC_ERROR_IF(v_fields.size() != 1, "FunctionalResidual nonlinear systems only supports 1 output residual");
+    SLIC_ERROR_IF(vjp_sensitivities.size() != fields.size(),
+                  "Invalid number of field sensitivities relative to the number of fields");
+    SLIC_ERROR_IF(v_fields.size() != 1, "FunctionalResidual nonlinear systems only supports 1 output residual");
 
     dt_ = dt;
 
-    // TODO (EBC): add in a future PR...
-    // std::vector<mfem::Vector*> test_par_gf({&v_fields[0]->gridFunction()});
-    // std::vector<mfem::Vector*> field_par_gf = getLVectors(fields);
-    // // field_par_gf.push_back(&v_fields[0]->gridFunction());
+    std::vector<mfem::Vector*> test_par_gf({&v_fields[0]->gridFunction()});
+    std::vector<mfem::Vector*> field_par_gf = getLVectors(fields);
 
-    // for (size_t input_col = 0; input_col < fields.size(); ++input_col) {
-    //   if (vjp_sensitivities[input_col] != nullptr) {
-    //     auto deriv_op = v_dot_weak_form_residual_.GetDerivative(input_col, test_par_gf, field_par_gf);
-    //     // do this entry by entry until assembly is supported
-    //     mfem::Vector direction(vjp_sensitivities[input_col]->Size());
-    //     direction = 0.0;
-    //     for (int i = 0; i < vjp_sensitivities[input_col]->Size(); ++i) {
-    //       direction[i] = 1.0;
-    //       mfem::Vector value(1);
-    //       deriv_op->Mult(direction, value);
-    //       (*vjp_sensitivities[input_col])[i] += value[0];
-    //       direction[i] = 0.0;
-    //     }
-    //   }
-    // }
+    for (size_t input_col = 0; input_col < fields.size(); ++input_col) {
+      if (vjp_sensitivities[input_col] != nullptr) {
+        auto deriv_op = weak_form_.GetDerivative(input_col, test_par_gf, field_par_gf);
+        deriv_op->MultTranspose(*v_fields[0], *vjp_sensitivities[input_col]);
+      }
+    }
   }
 
  protected:

src/smith/physics/tests/CMakeLists.txt

@@ -52,7 +52,8 @@ endif()
 
 if (SMITH_USE_DFEM)
     list(APPEND physics_serial_test_sources
-                test_dfem_explicit_dynamics.cpp)
+                test_dfem_explicit_dynamics.cpp
+                test_dfem_weak_form.cpp)
 endif()
 
 smith_add_tests(SOURCES       ${physics_serial_test_sources}