Fix indexing error in building of derivative operator

fem/dfem/doperator.hpp

@@ -143,6 +143,30 @@ public:
       prolongation_transpose(daction_l, result_t);
    };
 
+   void Assemble(Vector& result_t) const
+   {
+      auto l_width = GetVSize(direction);
+
+      daction_l.SetSize(l_width);
+      daction_l = 0.0;
+
+      direction_l.SetSize(l_width);
+      direction_l = 0.0;
+      for (const auto &f : derivative_actions)
+      {
+         for (int i{0}; i < direction_l.Size(); ++i)
+         {
+            direction_l[i] = 1.0;
+            mfem::Vector local_daction_l(daction_l_size);
+            local_daction_l = 0.0;
+            f(fields_e, direction_l, local_daction_l);
+            daction_l[i] += local_daction_l[0];
+            direction_l[i] = 0.0;
+         }
+      }
+      prolongation_transpose(daction_l, result_t);
+   }
+
    /// @brief Compute the transpose of the derivative operator on a given
    /// vector.
    ///
@@ -448,7 +472,7 @@ public:
 
       // Dummy
       Vector dir_l;
-      if (derivative_idx > s_l.size())
+      if (derivative_idx >= s_l.size())
       {
          dir_l = p_l[derivative_idx - s_l.size()];
       }
@@ -471,7 +495,7 @@ public:
                 sol_l,
                 par_l,
                 restriction_callback,
-                prolongation_transpose,
+                derivative_prolongation_transpose,
                 assemble_derivative_sparsematrix_callbacks[derivative_id],
                 assemble_derivative_hypreparmatrix_callbacks[derivative_id]);
    }
@@ -507,6 +531,7 @@ private:
    mutable Vector residual_e;
 
    std::function<void(Vector &, Vector &)> prolongation_transpose;
+   std::function<void(Vector &, Vector &)> derivative_prolongation_transpose;
    std::function<void(Vector &, Vector &)> output_restriction_transpose;
    restriction_callback_t restriction_callback;
 
@@ -727,7 +752,7 @@ void DifferentiableOperator::AddIntegrator(
    const int num_entities = GetNumEntities<entity_t>(mesh);
    const int num_qp = integration_rule.GetNPoints();
 
-   if constexpr (is_sum_fop<decltype(output_fop)>::value)
+   if constexpr (is_sum_fop<std::remove_const_t<decltype(output_fop)>>::value)
    {
       residual_l.SetSize(1);
       height = 1;
@@ -759,10 +784,13 @@ void DifferentiableOperator::AddIntegrator(
    auto input_dtq_maps = create_dtq_maps<entity_t>(inputs, dtq, input_to_field);
    auto output_dtq_maps = create_dtq_maps<entity_t>(outputs, dtq, output_to_field);
 
-   const int test_vdim = output_fop.vdim;
-   const int test_op_dim = output_fop.size_on_qp / output_fop.vdim;
+   int test_vdim = output_fop.vdim;
+   if constexpr (is_sum_fop<std::remove_const_t<decltype(output_fop)>>::value) {
+      test_vdim = 1;
+   }
+   const int test_op_dim = output_fop.size_on_qp / test_vdim;
    const int num_test_dof =
-      num_entities ? (output_e_size / output_fop.vdim / num_entities) : 0;
+      num_entities ? std::max(1, output_e_size / test_vdim / num_entities) : 0;
 
    auto ir_weights = Reshape(integration_rule.GetWeights().Read(), num_qp);
 
@@ -885,6 +913,17 @@ void DifferentiableOperator::AddIntegrator(
    // will fail.
    if constexpr (derivative_ids_t::size() != 0)
    {
+      if constexpr (is_sum_fop<std::remove_const_t<decltype(output_fop)>>::value)
+      {
+         // NOTE: the field ID in Value isn't used, so set it to anything
+         derivative_prolongation_transpose = get_prolongation_transpose(
+            fields[test_space_field_idx], Value<0>(), mesh.GetComm());
+      }
+      else
+      {
+         derivative_prolongation_transpose = prolongation_transpose;
+      }
+
       // Create the action of the derivatives
       for_constexpr([&, &or_transpose =
                         this->output_restriction_transpose](const std::size_t derivative_id)

fem/dfem/integrate.hpp

@@ -90,7 +90,7 @@ void map_quadrature_data_to_fields_impl(
    }
    else
    {
-      MFEM_ABORT("quadrature data mapping to field is not implemented for"
+      MFEM_ABORT_KERNEL("quadrature data mapping to field is not implemented for"
                  " this field descriptor");
    }
 }
@@ -306,8 +306,8 @@ void map_quadrature_data_to_fields_tensor_impl_2d(
    }
    else
    {
-      MFEM_ABORT("quadrature data mapping to field is not implemented for"
-                 " this field descriptor with sum factorization on tensor product elements");
+      MFEM_ABORT_KERNEL("quadrature data mapping to field is not implemented for"
+                        " this field descriptor with sum factorization on tensor product elements");
    }
 }
 
@@ -492,8 +492,8 @@ void map_quadrature_data_to_fields_tensor_impl_3d(
    }
    else
    {
-      MFEM_ABORT("quadrature data mapping to field is not implemented for"
-                 " this field descriptor with sum factorization on tensor product elements");
+      MFEM_ABORT_KERNEL("quadrature data mapping to field is not implemented for"
+                        " this field descriptor with sum factorization on tensor product elements");
    }
 }
 

fem/dfem/parameterspace.hpp

@@ -40,6 +40,9 @@ public:
    /// returns the local size of the space
    virtual int GetVSize() const = 0;
 
+   /// Number of elements
+   virtual int GetNE() const = 0;
+
    /// Get spatial dimension
    ///
    /// returns always 1.
@@ -74,6 +77,7 @@ public:
       return elem_restr.get();
    }
 
+
 protected:
    int vdim;
    DofToQuad dtq;
@@ -96,7 +100,7 @@ public:
    /// taken directly from the integration rule.
    UniformParameterSpace(Mesh &mesh, const IntegrationRule &ir, int vdim,
                          bool used_in_tensor_product = true) :
-      ParameterSpace(vdim)
+      ParameterSpace(vdim), ne(mesh.GetNE())
    {
       // Setup DofToQuad information
       dtq.nqpt = (int)floor(std::pow(ir.GetNPoints(), 1.0 / mesh.Dimension()) + 0.5);
@@ -122,12 +126,20 @@ public:
       return lsize;
    }
 
+   int GetNE() const override
+   {
+      return ne;
+   }
+
 private:
    /// T-vector size
    int tsize;
 
    /// L-vector size
    int lsize;
+
+   /// Number of elements
+   int ne;
 };
 
 class ParameterFunction : public Vector

fem/dfem/util.hpp

@@ -1040,6 +1040,29 @@ const Operator *get_restriction(const FieldDescriptor &f,
    return nullptr;
 }
 
+template <typename entity_t>
+inline int get_num_entities(const FieldDescriptor& f)
+{
+   if constexpr (std::is_same_v<entity_t, Entity::Element>) {
+      return std::visit(
+         [](auto&& arg) {
+            using T = std::decay_t<decltype(arg)>;
+            if constexpr (std::is_same_v<T, const FiniteElementSpace*> ||
+                          std::is_same_v<T, const ParFiniteElementSpace*>) {
+               return arg->GetNE();
+            } else if constexpr (std::is_same_v<T, const ParameterSpace*>) {
+               // TODO: implement GetNE() on ParameterSpace
+               return arg->GetNE();
+            } else {
+               static_assert(dfem::always_false<T>, "can't use get_num_entities on type");
+            }
+            return 0;  // Unreachable, but avoids compiler warning
+         },
+         f.data);
+  }
+  MFEM_ABORT("get_num_entities not implemented for Entity");
+}
+
 /// @brief Get a transpose restriction callback for a field descriptor.
 ///
 /// @param f the field descriptor.
@@ -1058,9 +1081,10 @@ get_restriction_transpose(
    {
       auto RT = [=](const Vector &v_e, Vector &v_l)
       {
-         v_l += v_e;
+         v_l += v_e.Sum();
       };
-      return std::make_tuple(RT, 1);
+      auto n_el = get_num_entities<entity_t>(f);
+      return std::make_tuple(RT, n_el);
    }
    else
    {
@@ -1187,7 +1211,9 @@ std::function<void(const Vector&, Vector&)> get_prolongation_transpose(
       {
          MFEM_ASSERT(y.Size() == 1, "output size doesn't match kernel description");
          real_t local_sum = r_local.Sum();
-         MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM, mpi_comm);
+         // NOTE: This doesn't work for a DerivativeOperator. You need to apply prolongation transpose on the field, 
+         // since the DOF IDs are different across ranks.
+         // MPI_Allreduce(&local_sum, y.GetData(), 1, MPI_DOUBLE, MPI_SUM, mpi_comm);
       };
       return PT;
    }