style

fem/dfem/backends/local_qf/derivative_assemble.hpp

@@ -16,8 +16,10 @@
 #include "kernels.hpp"
 #include "util.hpp"
 
+#include <algorithm>
 #include <array>
 #include <type_traits>
+#include <vector>
 
 namespace ker = mfem::kernels::internal;
 
@@ -592,10 +594,10 @@ class DerivativeAssemble
    const std::array<DofToQuadMap, n_outputs> output_dtq_maps;
    const std::array<bool, n_inputs> input_is_dependent;
    const size_t trial_field_uf;
-   const size_t test_field_uf;
-   const ParFiniteElementSpace *test_fes;
+   /// Column space of every row block. GetDerivative differentiates w.r.t. one
+   /// field, so there is exactly one trial space. Null when that field is not
+   /// an FE space, in which case nothing can be assembled.
    const ParFiniteElementSpace *trial_fes;
-   const int test_vdim;
    /// Per-output row geometry of the quadrature point cache. DerivativeSetup
    /// lays that cache out over every output FieldOperator, so reading it needs
    /// all of them.
@@ -603,7 +605,19 @@ class DerivativeAssemble
    const std::array<int, n_outputs> out_op_dim;
    const std::array<int, n_outputs> out_offsets;
    const int output_size_on_qp;
-   const int num_test_dof;
+   /// Output field ids, in order of first appearance among @a outputs. Each one
+   /// is a row block of the derivative: see @ref compute_group_field_ids.
+   const std::vector<int> group_field_ids;
+   /// Output FieldOperator -> row block it contributes to.
+   const std::array<int, n_outputs> out_group;
+   /// Row block -> position in ctx.outfds (-1 if the field is not an output of
+   /// the operator, which should not happen).
+   const std::vector<int> group_outfd_idx;
+   /// Row block -> test space, null for spaces without a basis.
+   const std::vector<const ParFiniteElementSpace *> group_fes;
+   const std::vector<int> group_test_vdim;
+   const std::vector<int> group_num_test_dof;
+   const std::vector<bool> group_assemblable;
    const int trial_vdim;
    const int trial_op_dim;
    const int num_trial_dof;
@@ -611,7 +625,36 @@ class DerivativeAssemble
    const int num_trial_dof_1d;
    const int total_trial_op_dim;
    mutable Vector inputs_trial_op_dim;
-   mutable Vector Ae_mem;
+   /// One element matrix bank per row block. Outputs on the same test field
+   /// share a bank and accumulate into it; outputs on different test fields are
+   /// different row blocks and need banks of their own, because their element
+   /// matrices have different shapes and are filled through different
+   /// ElementRestrictions.
+   mutable std::vector<Vector> group_Ae_mem;
+
+   /// @brief Distinct output field ids, in order of first appearance.
+   ///
+   /// Outputs sharing a field id form one row block: they are summed into a
+   /// single element matrix and produce a single assembled matrix. This is the
+   /// multi-output case, e.g. Outputs<Value<U>, Gradient<U>> giving mass plus
+   /// diffusion. Outputs on different field ids are separate row blocks of
+   ///
+   ///     dR/dU = [ dR_U/dU ; dR_Y/dU ]
+   ///
+   /// and are assembled into one matrix each.
+   static std::vector<int> compute_group_field_ids(const outputs_t &outs)
+   {
+      std::vector<int> ids;
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         const int fid = get<o>(outs).GetFieldId();
+         if (std::find(ids.begin(), ids.end(), fid) == ids.end())
+         {
+            ids.push_back(fid);
+         }
+      });
+      return ids;
+   }
 
 public:
    DerivativeAssemble() = delete;
@@ -658,29 +701,18 @@ public:
                       ctx_in.ir)),
    input_is_dependent(compute_input_is_dependent(inputs, derivative_id)),
    trial_field_uf(find_union_field_index(ctx_in, derivative_id)),
-   test_field_uf(
-      find_union_field_index(ctx_in, get<0>(outputs).GetFieldId())),
-   test_fes(
-      [&]
-   {
-      const auto *fes = std::get_if<const ParFiniteElementSpace *>(
-         &ctx_in.unionfds[test_field_uf].data);
-      MFEM_ASSERT(fes != nullptr && *fes != nullptr,
-                  "LocalQFBackend: test space is not a ParFiniteElementSpace");
-      return *fes;
-   }()),
    trial_fes(
-      [&]
+      [&]() -> const ParFiniteElementSpace *
    {
+      // Not every field is an FE space: a derivative w.r.t. a ParameterSpace or
+      // a QuadratureSpace has no basis to assemble columns against. That is
+      // only an error if assembly is actually requested, so keep it null here
+      // and report it in operator() instead of aborting registration.
+      if (trial_field_uf >= ctx_in.unionfds.size()) { return nullptr; }
       const auto *fes = std::get_if<const ParFiniteElementSpace *>(
          &ctx_in.unionfds[trial_field_uf].data);
-      MFEM_ASSERT(fes != nullptr && *fes != nullptr,
-                  "LocalQFBackend: trial space is not a ParFiniteElementSpace");
-      return *fes;
+      return fes ? *fes : nullptr;
    }()),
-   // All outputs are attached to the same test field, so vdim is common to
-   // them; only the operator dimension differs, and that lives in out_op_dim.
-   test_vdim(get<0>(outputs).vdim),
    out_vdim(get_vdim(outputs)),
    out_op_dim(compute_out_op_dim(outputs)),
    out_offsets(compute_out_offsets(out_vdim, out_op_dim)),
@@ -691,7 +723,96 @@ public:
       for_constexpr<n_outputs>([&](auto o) { s += get<o>(outputs).size_on_qp; });
       return s;
    }()),
-   num_test_dof(test_fes->GetFE(0)->GetDof()),
+   group_field_ids(compute_group_field_ids(outputs_in)),
+   out_group(
+      [&]
+   {
+      std::array<int, n_outputs> g {};
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         const int fid = get<o>(outputs_in).GetFieldId();
+         const auto &ids = group_field_ids;
+         g[o] = static_cast<int>(std::find(ids.begin(), ids.end(), fid)
+                                 - ids.begin());
+      });
+      return g;
+   }()),
+   group_outfd_idx(
+      [&]
+   {
+      std::vector<int> idx(group_field_ids.size(), -1);
+      for (size_t g = 0; g < group_field_ids.size(); g++)
+      {
+         for (size_t f = 0; f < ctx_in.outfds.size(); f++)
+         {
+            if (static_cast<int>(ctx_in.outfds[f].id) == group_field_ids[g])
+            {
+               idx[g] = static_cast<int>(f);
+               break;
+            }
+         }
+      }
+      return idx;
+   }()),
+   group_fes(
+      [&]
+   {
+      std::vector<const ParFiniteElementSpace *> v(group_field_ids.size(),
+                                                   nullptr);
+      for (size_t g = 0; g < v.size(); g++)
+      {
+         if (group_outfd_idx[g] < 0) { continue; }
+         const auto *fes = std::get_if<const ParFiniteElementSpace *>(
+                              &ctx_in.outfds[group_outfd_idx[g]].data);
+         v[g] = fes ? *fes : nullptr;
+      }
+      return v;
+   }()),
+   group_test_vdim(
+      [&]
+   {
+      // Outputs in a group share a field, hence a vdim; only the operator
+      // dimension differs between them, and that lives in out_op_dim.
+      std::vector<int> v(group_field_ids.size(), 0);
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         v[out_group[o]] = get<o>(outputs_in).vdim;
+      });
+      return v;
+   }()),
+   group_num_test_dof(
+      [&]
+   {
+      std::vector<int> v(group_field_ids.size(), 0);
+      for (size_t g = 0; g < v.size(); g++)
+      {
+         if (group_fes[g] == nullptr) { continue; }
+         v[g] = group_fes[g]->GetFE(0)->GetDof();
+      }
+      return v;
+   }()),
+   group_assemblable(
+      [&]
+   {
+      // A row block can only be assembled when both its trial and test fields are
+      // ParFiniteElementSpaces. Quadrature and parameter spaces have no element
+      // basis or ElementRestriction through which to form a SparseMatrix.
+      //
+      // Identity outputs are also excluded for now as it has no supported mapping from
+      // its pointwise quadrature rows into that FE row space.
+      std::vector<bool> v(group_field_ids.size(), false);
+      if (trial_fes == nullptr) { return v; }
+      for (size_t g = 0; g < v.size(); g++) { v[g] = group_fes[g] != nullptr; }
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         using output_fop_t = std::decay_t<decltype(get<o>(outputs_in))>;
+         if constexpr (is_identity_fop_v<output_fop_t>)
+         {
+            v[out_group[o]] = false;
+         }
+      });
+      return v;
+   }()),
    trial_vdim(compute_trial_vdim(inputs, derivative_id)), trial_op_dim(
       [&]
    {
@@ -705,7 +826,7 @@ public:
       });
       return top;
    }()),
-   num_trial_dof(trial_fes->GetFE(0)->GetDof()),
+   num_trial_dof(trial_fes ? trial_fes->GetFE(0)->GetDof() : 0),
    dim(ctx_in.mesh.Dimension()), ne(ctx_in.nentities),
    nq(ctx_in.ir.GetNPoints()), q1d(tensor_1d_size(nq, dim)),
    num_trial_dof_1d(tensor_1d_size(num_trial_dof, dim)), total_trial_op_dim(
@@ -716,14 +837,12 @@ public:
       return compute_total_trial_op_dim(
          inputs, input_is_dependent, in_qp_sizes);
    }()),
-   inputs_trial_op_dim(), Ae_mem()
+   inputs_trial_op_dim(), group_Ae_mem()
    {
       MFEM_ASSERT(ctx.unionfds.size() == nfields,
                   "LocalQFBackend: unionfds size mismatch");
       MFEM_ASSERT(trial_field_uf != SIZE_MAX,
                   "DerivativeAssemble: trial field not found in unionfds");
-      MFEM_ASSERT(test_field_uf != SIZE_MAX,
-                  "DerivativeAssemble: test field not found in unionfds");
 
       MFEM_ASSERT(trial_vdim > 0,
                   "LocalQFBackend: could not determine trial vdim");
@@ -741,79 +860,114 @@ public:
             : 0;
       });
 
-      const int elem_mat_size =
-         num_test_dof * test_vdim * num_trial_dof * trial_vdim;
-      Ae_mem.SetSize(elem_mat_size * ne, Device::GetDeviceMemoryType());
-      Ae_mem.UseDevice(true);
-      Ae_mem = 0.0;
+      group_Ae_mem.resize(group_field_ids.size());
+      for (size_t g = 0; g < group_Ae_mem.size(); g++)
+      {
+         if (!group_assemblable[g]) { continue; }
+         const int elem_mat_size = group_num_test_dof[g] * group_test_vdim[g] *
+                                   num_trial_dof * trial_vdim;
+         group_Ae_mem[g].SetSize(elem_mat_size * ne,
+                                 Device::GetDeviceMemoryType());
+         group_Ae_mem[g].UseDevice(true);
+         group_Ae_mem[g] = 0.0;
+      }
    }
 
-   void operator()(SparseMatrix *&A) const
+   /// @brief Assemble one SparseMatrix per assemblable output field.
+   ///
+   /// @a A is indexed by position in ctx.outfds, so A[f] receives the row block
+   /// belonging to output field f. Slots this integrator writes nothing
+   /// assemblable to -- quadrature or parameter spaces, or fields it does not
+   /// touch at all -- are left alone, which lets several integrators fill
+   /// different row blocks of the same vector. Ownership of the matrices passes
+   /// to the caller.
+   void operator()(std::vector<SparseMatrix *> &A) const
    {
-      // Every output is contracted into one element matrix Ae, sized from the
-      // test space of get<0>(outputs), and filled through a single test
-      // ElementRestriction.
-      //
-      // WIP:
-      // This takes care of single-field, multiple-outputs case.
-      // For a multiple fields case, outputs on a second field would need a second
-      // matrix -- the derivative then eould be a block column with one row block per
-      // test space.
-      //
-      // For now we just add a check that all outputs are attached to the same test field, and abort if not.
-      for_constexpr<n_outputs>([&](auto o)
-      {
-         MFEM_VERIFY(get<o>(outputs).GetFieldId() == get<0>(outputs).GetFieldId(),
-                     "DerivativeAssemble: every output FieldOperator has to be "
-                     "attached to the same test field; assembling outputs that "
-                     "span several fields is not supported");
-      });
-
       if (ctx.attr.Size() == 0) { return; }
 
+      const bool any_assemblable =
+         std::find(group_assemblable.begin(), group_assemblable.end(), true) !=
+         group_assemblable.end();
+      if (!any_assemblable) { return; }
+
       if (!(use_sum_factorization && (dim == 2 || dim == 3)))
       {
          MFEM_ABORT("DerivativeAssemble optimized path is implemented "
                     "for tensor-product 2D/3D elements only");
       }
 
-      DerivativeAssembleHO::Run(dim,
-                                q1d,
-                                ctx,
-                                qp_cache,
-                                Ae_mem,
-                                inputs,
-                                outputs,
-                                input_dtq_maps,
-                                output_dtq_maps,
-                                out_vdim,
-                                out_op_dim,
-                                out_offsets,
-                                output_size_on_qp,
-                                inputs_trial_op_dim,
-                                test_vdim,
-                                num_test_dof,
-                                num_trial_dof,
-                                num_trial_dof_1d,
-                                trial_vdim,
-                                total_trial_op_dim,
-                                nq,
-                                ne,
-                                q1d,
-                                dim);
+      MFEM_VERIFY(trial_fes != nullptr,
+                  "DerivativeAssemble: the differentiated field is not a "
+                  "ParFiniteElementSpace, so the columns of the derivative "
+                  "have no basis to be assembled against");
 
-      A = new SparseMatrix;
-      A->OverrideSize(test_fes->GetVSize(), trial_fes->GetVSize());
-
-      const auto *test_restr = dynamic_cast<const ElementRestriction *>(
-                                  test_fes->GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC));
       const auto *trial_restr = dynamic_cast<const ElementRestriction *>(
-                                   trial_fes->GetElementRestriction(ElementDofOrdering::LEXICOGRAPHIC));
-      MFEM_VERIFY(test_restr != nullptr && trial_restr != nullptr,
+                                   trial_fes->GetElementRestriction(
+                                      ElementDofOrdering::LEXICOGRAPHIC));
+      MFEM_VERIFY(trial_restr != nullptr,
                   "DerivativeAssemble SparseMatrix assembly requires "
                   "H1/conforming ElementRestriction spaces");
 
-      test_restr->FillSparseMatrix(Ae_mem, *A, *trial_restr);
+      if (A.size() < ctx.outfds.size())
+      {
+         A.resize(ctx.outfds.size(), nullptr);
+      }
+
+      for (size_t g = 0; g < group_Ae_mem.size(); g++)
+      {
+         if (!group_assemblable[g]) { continue; }
+
+         // The kernel accumulates into Ae, so the bank has to start clean on
+         // every call: otherwise assembling twice (the next Newton step, say)
+         // would double every entry.
+         group_Ae_mem[g] = 0.0;
+
+         DerivativeAssembleHO::Run(dim,
+                                   q1d,
+                                   ctx,
+                                   qp_cache,
+                                   group_Ae_mem[g],
+                                   inputs,
+                                   outputs,
+                                   input_dtq_maps,
+                                   output_dtq_maps,
+                                   out_vdim,
+                                   out_op_dim,
+                                   out_offsets,
+                                   out_group,
+                                   static_cast<int>(g),
+                                   output_size_on_qp,
+                                   inputs_trial_op_dim,
+                                   group_test_vdim[g],
+                                   group_num_test_dof[g],
+                                   num_trial_dof,
+                                   num_trial_dof_1d,
+                                   trial_vdim,
+                                   total_trial_op_dim,
+                                   nq,
+                                   ne,
+                                   q1d,
+                                   dim);
+
+         const ParFiniteElementSpace *test_fes = group_fes[g];
+         const auto *test_restr = dynamic_cast<const ElementRestriction *>(
+                                     test_fes->GetElementRestriction(
+                                        ElementDofOrdering::LEXICOGRAPHIC));
+         MFEM_VERIFY(test_restr != nullptr,
+                     "DerivativeAssemble SparseMatrix assembly requires "
+                     "H1/conforming ElementRestriction spaces");
+
+         const int f = group_outfd_idx[g];
+         MFEM_VERIFY(A[f] == nullptr,
+                     "DerivativeAssemble: output field already carries an "
+                     "assembled matrix; two integrators contributing to the "
+                     "same row block cannot be assembled into one matrix");
+
+         auto *M = new SparseMatrix;
+         M->OverrideSize(test_fes->GetVSize(), trial_fes->GetVSize());
+         test_restr->FillSparseMatrix(group_Ae_mem[g], *M, *trial_restr);
+         A[f] = M;
+      }
    }
 
    template<typename backend_t = LocalQFHOBackend<3>, int T_Q1D = 0>
@@ -828,6 +982,8 @@ public:
       const std::array<int, n_outputs> &out_vdim,
       const std::array<int, n_outputs> &out_op_dim,
       const std::array<int, n_outputs> &out_offsets,
+      const std::array<int, n_outputs> &out_group,
+      const int group,
       const int output_size_on_qp,
       const Vector &inputs_trial_op_dim,
       const int test_vdim,
@@ -853,6 +1009,7 @@ public:
                   "DerivativeAssemble: nq exceeds backend quadrature capacity");
       for_constexpr<n_outputs>([&](auto o)
       {
+         if (out_group[o] != group) { return; }
          MFEM_VERIFY(out_op_dim[o] <= DIM,
                      "DerivativeAssemble: test_op_dim exceeds spatial DIM");
       });
@@ -896,14 +1053,25 @@ public:
 
          // Each output contributes its own rows of the cache, contracted
          // against its own test basis operation; map_quadrature_data_to_fields
-         // accumulates, so the element matrix is the sum over outputs for the
-         // same field.
+         // accumulates, so the element matrix is the sum over the outputs that
+         // belong to this row block. Outputs on other test fields are skipped
+         // here and picked up by their own launch, which has its own Ae.
+         //
+         // A launch only ever happens for an assemblable row block, and such a
+         // block contains no Identity output, so the is_identity_fop_v test
+         // below never rejects anything at run time. It is there to keep
+         // assemble_element_mat_sumfact from being instantiated for a fop it
+         // cannot handle.
          for_constexpr<n_outputs>([&](auto o)
          {
             using output_fop_t = std::decay_t<decltype(get<o>(outputs))>;
 
             if constexpr (!is_identity_fop_v<output_fop_t>)
             {
+               // Uniform across the thread block, so returning early here
+               // cannot desynchronise the barriers below.
+               if (out_group[o] != group) { return; }
+
                // The outputs share fhat_storage, so one has to be done with it
                // before the next zeroes it.
                MFEM_SYNC_THREAD;

fem/dfem/backends/local_qf/derivative_assemble_diagonal.hpp

@@ -15,12 +15,18 @@
 #include "kernels.hpp"
 #include "util.hpp"
 
+#include <algorithm>
 #include <array>
+#include <vector>
 
 namespace mfem::future::LocalQFImpl
 {
 
-// Assemble diagonal of cached Jacobian (square trial == test, tensor 2D/3D)
+// Assemble the diagonal of one row block of a cached Jacobian (tensor 2D/3D).
+//
+// The derivative is a block column, one row block per output field. Only a
+// block whose test space is the trial space is square, and only a square block
+// has a diagonal at all, so the row block is chosen per call and checked.
 
 template<int derivative_id,
          typename qfunc_t,
@@ -47,22 +53,47 @@ class DerivativeAssembleDiagonal
    const std::array<DofToQuadMap, n_outputs> output_dtq_maps;
    const std::array<bool, n_inputs> input_is_dependent;
    const size_t trial_field_uf;
-   const size_t test_field_uf;
-   const bool is_square;
-   const int test_vdim;
+   /// Column space of every row block; null when the differentiated field is
+   /// not an FE space, in which case no block has a diagonal.
+   const ParFiniteElementSpace *trial_fes;
    const std::array<int, n_outputs> out_vdim;
    const std::array<int, n_outputs> out_op_dim;
    const std::array<int, n_outputs> out_offsets;
    const int output_size_on_qp;
-   const int num_test_dof;
-   const int num_test_dof_1d;
+   /// Row blocks: the distinct output field ids, in order of first appearance.
+   /// Outputs sharing a field id are summed into one diagonal, which is how
+   /// Value<U> + Gradient<U> becomes mass plus diffusion.
+   const std::vector<int> group_field_ids;
+   const std::array<int, n_outputs> out_group;
+   const std::vector<const ParFiniteElementSpace *> group_fes;
+   const std::vector<int> group_test_vdim;
+   const std::vector<int> group_num_test_dof;
+   const std::vector<int> group_num_test_dof_1d;
+   /// Whether a row block has a diagonal: its test space has to be an FE space
+   /// and has to *be* the trial space, and no output on it may be an Identity.
+   const std::vector<bool> group_has_diagonal;
    const int trial_vdim;
    const int total_trial_op_dim;
    const int num_trial_dof_1d;
    const int residual_size_on_qp;
    const int dim, ne, nq, q1d;
    const std::array<int, n_inputs> inputs_trial_op_dim;
-   mutable Vector Ye_mem;
+   mutable std::vector<Vector> group_Ye_mem;
+
+   /// Distinct output field ids, in order of first appearance.
+   static std::vector<int> compute_group_field_ids(const outputs_t &outs)
+   {
+      std::vector<int> ids;
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         const int fid = get<o>(outs).GetFieldId();
+         if (std::find(ids.begin(), ids.end(), fid) == ids.end())
+         {
+            ids.push_back(fid);
+         }
+      });
+      return ids;
+   }
 
 public:
    DerivativeAssembleDiagonal() = delete;
@@ -109,19 +140,14 @@ public:
                       ctx_in.ir)),
    input_is_dependent(compute_input_is_dependent(inputs, derivative_id)),
    trial_field_uf(find_union_field_index(ctx_in, derivative_id)),
-   test_field_uf(
-      find_union_field_index(ctx_in, get<0>(outputs).GetFieldId())),
-   is_square(
-      [&]
+   trial_fes(
+      [&]() -> const ParFiniteElementSpace *
    {
-      const auto *test_fes = std::get_if<const ParFiniteElementSpace *>(
-         &ctx_in.unionfds[test_field_uf].data);
-      const auto *trial_fes = std::get_if<const ParFiniteElementSpace *>(
+      if (trial_field_uf >= ctx_in.unionfds.size()) { return nullptr; }
+      const auto *fes = std::get_if<const ParFiniteElementSpace *>(
          &ctx_in.unionfds[trial_field_uf].data);
-      return test_fes && trial_fes && *test_fes && *trial_fes &&
-      (*test_fes == *trial_fes);
+      return fes ? *fes : nullptr;
    }()),
-   test_vdim(get<0>(outputs).vdim),
    out_vdim(get_vdim(outputs_in)),
    out_op_dim(compute_out_op_dim(outputs_in)),
    out_offsets(compute_out_offsets(out_vdim, out_op_dim)),
@@ -132,16 +158,96 @@ public:
       for_constexpr<n_outputs>([&](auto o)
       { s += get<o>(outputs_in).size_on_qp; });
       return s;
-   }()), num_test_dof(
+   }()),
+   group_field_ids(compute_group_field_ids(outputs_in)),
+   out_group(
       [&]
    {
-      const auto *test_fes = std::get_if<const ParFiniteElementSpace *>(
-         &ctx_in.unionfds[test_field_uf].data);
-      MFEM_ASSERT(test_fes != nullptr && *test_fes != nullptr,
-                  "LocalQFBackend: test space is not a ParFiniteElementSpace");
-      return (*test_fes)->GetFE(0)->GetDof();
+      std::array<int, n_outputs> g {};
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         const int fid = get<o>(outputs_in).GetFieldId();
+         const auto &ids = group_field_ids;
+         g[o] = static_cast<int>(std::find(ids.begin(), ids.end(), fid)
+                                 - ids.begin());
+      });
+      return g;
+   }()),
+   group_fes(
+      [&]
+   {
+      std::vector<const ParFiniteElementSpace *> v(group_field_ids.size(),
+                                                   nullptr);
+      for (size_t g = 0; g < v.size(); g++)
+      {
+         const size_t uf = find_union_field_index(ctx_in, group_field_ids[g]);
+         if (uf >= ctx_in.unionfds.size()) { continue; }
+         const auto *fes = std::get_if<const ParFiniteElementSpace *>(
+                              &ctx_in.unionfds[uf].data);
+         v[g] = fes ? *fes : nullptr;
+      }
+      return v;
+   }()),
+   group_test_vdim(
+      [&]
+   {
+      std::vector<int> v(group_field_ids.size(), 0);
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         v[out_group[o]] = get<o>(outputs_in).vdim;
+      });
+      return v;
+   }()),
+   group_num_test_dof(
+      [&]
+   {
+      std::vector<int> v(group_field_ids.size(), 0);
+      for (size_t g = 0; g < v.size(); g++)
+      {
+         if (group_fes[g] == nullptr) { continue; }
+         v[g] = group_fes[g]->GetFE(0)->GetDof();
+      }
+      return v;
+   }()),
+   group_num_test_dof_1d(
+      [&]
+   {
+      std::vector<int> v(group_field_ids.size(), 0);
+      for (size_t g = 0; g < v.size(); g++)
+      {
+         if (group_num_test_dof[g] > 0)
+         {
+            v[g] = tensor_1d_size(group_num_test_dof[g],
+                                  ctx_in.mesh.Dimension());
+         }
+      }
+      return v;
+   }()),
+   group_has_diagonal(
+      [&]
+   {
+      // A diagonal needs row space == column space, so only a row block on the
+      // trial space qualifies. Squareness alone cannot pick a block when
+      // several output fields share that space, which is why the caller names
+      // the row. Identity outputs are quadrature point data and are excluded
+      // for the same reason as in DerivativeAssemble: they cannot be
+      // contracted, and every output on a field lands in the same block.
+      std::vector<bool> v(group_field_ids.size(), false);
+      if (trial_fes == nullptr) { return v; }
+      for (size_t g = 0; g < v.size(); g++)
+      {
+         v[g] = (group_fes[g] != nullptr) && (group_fes[g] == trial_fes);
+      }
+      for_constexpr<n_outputs>([&](auto o)
+      {
+         using output_fop_t = std::decay_t<decltype(get<o>(outputs_in))>;
+         if constexpr (is_identity_fop_v<output_fop_t>)
+         {
+            v[out_group[o]] = false;
+         }
+      });
+      return v;
    }()),
-   num_test_dof_1d(tensor_1d_size(num_test_dof, ctx_in.mesh.Dimension())),
    trial_vdim(compute_trial_vdim(inputs, derivative_id)), total_trial_op_dim(
       [&]
    {
@@ -151,15 +257,9 @@ public:
          inputs, input_is_dependent, input_size_on_qp);
    }()),
    num_trial_dof_1d(
-      [&]
-   {
-      const auto *trial_fes = std::get_if<const ParFiniteElementSpace *>(
-         &ctx_in.unionfds[trial_field_uf].data);
-      MFEM_ASSERT(trial_fes != nullptr && *trial_fes != nullptr,
-                  "LocalQFBackend: trial space is not a ParFiniteElementSpace");
-      const int num_trial_dof = (*trial_fes)->GetFE(0)->GetDof();
-      return tensor_1d_size(num_trial_dof, ctx_in.mesh.Dimension());
-   }()),
+      trial_fes ? tensor_1d_size(trial_fes->GetFE(0)->GetDof(),
+                                 ctx_in.mesh.Dimension())
+      : 0),
    residual_size_on_qp(output_size_on_qp * trial_vdim * total_trial_op_dim),
    dim(ctx_in.mesh.Dimension()), ne(ctx_in.nentities),
    nq(ctx_in.ir.GetNPoints()), q1d(tensor_1d_size(nq, dim)),
@@ -175,53 +275,62 @@ public:
       });
       return itod;
    }()),
-   Ye_mem()
+   group_Ye_mem()
    {
       MFEM_ASSERT(ctx.unionfds.size() == nfields,
                   "LocalQFBackend: unionfds size mismatch");
       MFEM_ASSERT(
          trial_field_uf != SIZE_MAX,
          "DerivativeAssembleDiagonal: trial field not found in unionfds");
-      MFEM_ASSERT(
-         test_field_uf != SIZE_MAX,
-         "DerivativeAssembleDiagonal: test field not found in unionfds");
       MFEM_ASSERT(trial_vdim > 0,
                   "LocalQFBackend: could not determine trial vdim");
       MFEM_ASSERT(total_trial_op_dim > 0,
                   "LocalQFBackend: no dependent inputs found");
       for_constexpr<n_outputs>([&](auto o)
       {
-         MFEM_CONTRACT_VAR(o);
-         MFEM_ASSERT(out_vdim[o] == test_vdim,
-                     "DerivativeAssembleDiagonal: all outputs must share the "
-                     "test field vdim");
+         MFEM_ASSERT(out_vdim[o] == group_test_vdim[out_group[o]],
+                     "DerivativeAssembleDiagonal: outputs on one field must "
+                     "share its vdim");
       });
 
-      if (is_square)
+      group_Ye_mem.resize(group_field_ids.size());
+      for (size_t g = 0; g < group_Ye_mem.size(); g++)
       {
-         Ye_mem.SetSize(num_test_dof * test_vdim * ne);
-         Ye_mem.UseDevice(true);
+         if (!group_has_diagonal[g]) { continue; }
+         group_Ye_mem[g].SetSize(group_num_test_dof[g] * group_test_vdim[g] *
+                                 ne);
+         group_Ye_mem[g].UseDevice(true);
       }
    }
 
+   /// Index of the row block for output field @a field_id, or -1.
+   int FindGroup(int field_id) const
+   {
+      const auto &ids = group_field_ids;
+      const auto it = std::find(ids.begin(), ids.end(), field_id);
+      return (it == ids.end()) ? -1 : static_cast<int>(it - ids.begin());
+   }
+
    template<typename Backend>
-   void run_kernels() const
+   void run_kernels(const int g) const
    {
       Backend::Run(dim,
                    q1d,
                    ctx,
                    qp_cache,
-                   Ye_mem,
+                   group_Ye_mem[g],
                    inputs,
                    outputs,
                    output_dtq_maps,
                    input_dtq_maps,
-                   test_vdim,
+                   out_group,
+                   g,
+                   group_test_vdim[g],
                    out_op_dim,
                    out_offsets,
                    output_size_on_qp,
-                   num_test_dof,
-                   num_test_dof_1d,
+                   group_num_test_dof[g],
+                   group_num_test_dof_1d[g],
                    trial_vdim,
                    total_trial_op_dim,
                    residual_size_on_qp,
@@ -232,9 +341,14 @@ public:
                    dim);
    }
 
-   void operator()(Vector &diag_e) const
+   /// Add this integrator's contribution to the diagonal of the row block of
+   /// output field @a out_field_id. Adds nothing if the integrator writes no
+   /// square, basis-backed block for that field; the caller is responsible for
+   /// rejecting a row that no integrator can serve.
+   void operator()(const int out_field_id, Vector &diag_e) const
    {
-      if (!is_square) { return; }
+      const int g = FindGroup(out_field_id);
+      if (g < 0 || !group_has_diagonal[g]) { return; }
       if (ctx.attr.Size() == 0) { return; }
 
       if (!(use_sum_factorization && (dim == 2 || dim == 3)))
@@ -242,27 +356,28 @@ public:
          MFEM_ABORT("DerivativeAssembleDiagonal optimized path is implemented "
                     "for tensor-product 2D/3D elements only");
       }
-      MFEM_VERIFY(num_test_dof_1d == num_trial_dof_1d,
+      MFEM_VERIFY(group_num_test_dof_1d[g] == num_trial_dof_1d,
                   "DerivativeAssembleDiagonal requires matching tensor dofs");
-      MFEM_VERIFY(num_test_dof_1d <= DeviceDofQuadLimits::Get().MAX_D1D, "");
-      MFEM_VERIFY(q1d <= DeviceDofQuadLimits::Get().MAX_Q1D, "");
+      const auto &limits = DeviceDofQuadLimits::Get();
+      MFEM_VERIFY(group_num_test_dof_1d[g] <= limits.MAX_D1D, "");
+      MFEM_VERIFY(q1d <= limits.MAX_Q1D, "");
 
-      Ye_mem = 0.0;
+      group_Ye_mem[g] = 0.0;
 
       if (q1d <= LocalQFLOBackendMQ1())
       {
-         run_kernels<DerivativeAssembleDiagonalLO>();
+         run_kernels<DerivativeAssembleDiagonalLO>(g);
       }
       else if (q1d <= LocalQFHOBackendMQ1())
       {
-         run_kernels<DerivativeAssembleDiagonalHO>();
+         run_kernels<DerivativeAssembleDiagonalHO>(g);
       }
       else
       {
          MFEM_ABORT("Unsupported quadrature order for LocalQF backend");
       }
 
-      diag_e += Ye_mem;
+      diag_e += group_Ye_mem[g];
    }
 
    template<typename backend_t = LocalQFLOBackend<3>, int T_Q1D = 0>
@@ -274,6 +389,8 @@ public:
       const outputs_t &outputs,
       const std::array<DofToQuadMap, n_outputs> &output_dtq_maps,
       const std::array<DofToQuadMap, n_inputs> &input_dtq_maps,
+      const std::array<int, n_outputs> &out_group,
+      const int row_group,
       const int test_vdim,
       const std::array<int, n_outputs> &out_op_dim,
       const std::array<int, n_outputs> &out_offsets,
@@ -339,85 +456,92 @@ public:
             }
             MFEM_SYNC_THREAD;
 
-            // Accumulate every (output o, test op k, dependent input s,
-            // trial op m) block of the cached Jacobian into the diagonal via
-            // the backend driver.
+            // Accumulate every output belonging to the requested row block.
+            // This sums multiple contributions, such as Value<U> +
+            // Gradient<U>, while skipping outputs on the other row blocks.
+            // The row is a run time choice, so unlike the field id it cannot
+            // gate the instantiation; is_identity_fop_v still does, since
+            // eval_test has no meaning for quadrature point data.
             for_constexpr<n_outputs>([&](auto o)
             {
                using test_fop_t = std::decay_t<decltype(get<o>(outputs))>;
-               const auto &out_dtq = output_dtq_maps[o];
-               const int test_op_dim = out_op_dim[static_cast<int>(o)];
-
-               // Test-basis factor along a spatial axis
-               const auto eval_test =
-                  [&](const int k, const int axis, const int q, const int d)
+               if constexpr (!is_identity_fop_v<test_fop_t>)
                {
-                  const auto &B = out_dtq.B;
-                  const auto &G = out_dtq.G;
-                  if constexpr (is_value_fop<test_fop_t>::value)
-                  {
-                     return (k == 0) ? B(q, 0, d) : 0.0;
-                  }
-                  else if constexpr (is_gradient_fop<test_fop_t>::value)
-                  {
-                     return (k == axis) ? G(q, 0, d) : B(q, 0, d);
-                  }
-                  else
-                  {
-                     return 0.0;
-                  }
-               };
+                  if (out_group[static_cast<int>(o)] != row_group) { return; }
+                  const auto &out_dtq = output_dtq_maps[o];
+                  const int test_op_dim = out_op_dim[static_cast<int>(o)];
 
-               for (int k = 0; k < test_op_dim; k++)
-               {
-                  const int row =
-                     out_offsets[static_cast<int>(o)] + vd * test_op_dim + k;
-                  int m_offset = 0;
-                  for_constexpr<n_inputs>([&](auto s)
+                  // Test-basis factor along a spatial axis
+                  const auto eval_test =
+                     [&](const int k, const int axis, const int q, const int d)
                   {
-                     using fop_t = std::decay_t<decltype(get<s>(inputs))>;
-                     const int trial_op_dim =
-                        inputs_trial_op_dim[static_cast<int>(s)];
-                     if (trial_op_dim == 0) { return; }
-
-                     const auto &in_dtq = input_dtq_maps[s];
-                     const auto eval_input =
-                        [&](const int m, const int axis, const int q,
-                            const int d)
+                     const auto &B = out_dtq.B;
+                     const auto &G = out_dtq.G;
+                     if constexpr (is_value_fop<test_fop_t>::value)
                      {
-                        if constexpr (is_value_fop<fop_t>::value)
-                        {
-                           return (m == 0) ? in_dtq.B(q, 0, d) : 0.0;
-                        }
-                        else if constexpr (is_gradient_fop<fop_t>::value)
-                        {
-                           return (m == axis) ? in_dtq.G(q, 0, d)
-                                  : in_dtq.B(q, 0, d);
-                        }
-                        else
-                        {
-                           return 0.0;
-                        }
-                     };
-
-                     for (int m = 0; m < trial_op_dim; m++)
-                     {
-                        const int col = m_offset + m;
-                        backend_t::DiagContract(
-                           s_diag,
-                           num_test_dof_1d,
-                           q1d,
-                           nz_dof,
-                           [&](int axis, int q, int d)
-                        { return eval_test(k, axis, q, d); },
-                        [&](int axis, int q, int d)
-                        { return eval_input(m, axis, q, d); },
-                        [&](int q) { return qpdc(q, col, vd, row); },
-                        [&](int dx, int dy, int dz, real_t u)
-                        { Y(dx, dy, dz) += u; });
+                        return (k == 0) ? B(q, 0, d) : 0.0;
                      }
-                     m_offset += trial_op_dim;
-                  });
+                     else if constexpr (is_gradient_fop<test_fop_t>::value)
+                     {
+                        return (k == axis) ? G(q, 0, d) : B(q, 0, d);
+                     }
+                     else
+                     {
+                        return 0.0;
+                     }
+                  };
+
+                  for (int k = 0; k < test_op_dim; k++)
+                  {
+                     const int row = out_offsets[static_cast<int>(o)] +
+                                     vd * test_op_dim + k;
+                     int m_offset = 0;
+                     for_constexpr<n_inputs>([&](auto s)
+                     {
+                        using fop_t = std::decay_t<decltype(get<s>(inputs))>;
+                        const int trial_op_dim =
+                           inputs_trial_op_dim[static_cast<int>(s)];
+                        if (trial_op_dim == 0) { return; }
+
+                        const auto &in_dtq = input_dtq_maps[s];
+                        const auto eval_input =
+                           [&](const int m, const int axis, const int q,
+                               const int d)
+                        {
+                           if constexpr (is_value_fop<fop_t>::value)
+                           {
+                              return (m == 0) ? in_dtq.B(q, 0, d) : 0.0;
+                           }
+                           else if constexpr (is_gradient_fop<fop_t>::value)
+                           {
+                              return (m == axis) ? in_dtq.G(q, 0, d)
+                                     : in_dtq.B(q, 0, d);
+                           }
+                           else
+                           {
+                              return 0.0;
+                           }
+                        };
+
+                        for (int m = 0; m < trial_op_dim; m++)
+                        {
+                           const int col = m_offset + m;
+                           backend_t::DiagContract(
+                              s_diag,
+                              num_test_dof_1d,
+                              q1d,
+                              nz_dof,
+                              [&](int axis, int q, int d)
+                           { return eval_test(k, axis, q, d); },
+                           [&](int axis, int q, int d)
+                           { return eval_input(m, axis, q, d); },
+                           [&](int q) { return qpdc(q, col, vd, row); },
+                           [&](int dx, int dy, int dz, real_t u)
+                           { Y(dx, dy, dz) += u; });
+                        }
+                        m_offset += trial_op_dim;
+                     });
+                  }
                }
             });
          }

fem/dfem/backends/local_qf/derivative_setup.hpp

@@ -384,6 +384,8 @@ public:
                            for_constexpr<n_outputs>([&](auto oc)
                            {
                               constexpr size_t o = oc.value, ao = n_inputs + o;
+                              using out_fop_t =
+                                 std::decay_t<tuple_element_t<o, outputs_t>>;
                               const auto &tangent = get<ao>(shadow_args);
                               const int tv = out_vdim[o], to = out_op_dim[o];
                               for (int i = 0; i < tv; i++)
@@ -394,8 +396,26 @@ public:
                                     const int cache_idx =
                                        row * trial_vdim * total_trial_op_dim +
                                        j * total_trial_op_dim + col_m;
-                                    cache_tensor(q, cache_idx, e) =
-                                       qf_value_at(tangent, i, k);
+                                    // An Identity output is flat quadrature
+                                    // point data: its FieldOperator vdim counts
+                                    // components, not rows of the q-function
+                                    // argument's shape. The two index form
+                                    // assumes vdim == extents[0] and would run
+                                    // off the end of, say, a
+                                    // tensor<real_t, DIM, DIM> bound to a vdim
+                                    // DIM*DIM space, so read it flat with the
+                                    // same column major packing that
+                                    // identity_qp_write_value writes.
+                                    if constexpr (is_identity_fop_v<out_fop_t>)
+                                    {
+                                       cache_tensor(q, cache_idx, e) =
+                                          qf_flat_value(tangent, i * to + k);
+                                    }
+                                    else
+                                    {
+                                       cache_tensor(q, cache_idx, e) =
+                                          qf_value_at(tangent, i, k);
+                                    }
                                  }
                               }
                            });

fem/dfem/doperator.hpp

@@ -109,19 +109,29 @@ struct derivative_action_t
    std::function<void *()> qfunc_shadow;
 };
 
-/// @brief Type alias for a function that assembles the SparseMatrix of a
-/// derivative operator
+/// @brief Type alias for a function that assembles the SparseMatrix row blocks
+/// of a derivative operator
+///
+/// The derivative of a residual with several output (test) fields is a block
+/// column, one row block per output field, so the callback fills a vector
+/// indexed by position in the operator's output FieldDescriptors rather than a
+/// single matrix.
 using assemble_derivative_sparsematrix_callback_t =
-   std::function<void(SparseMatrix *&)>;
+   std::function<void(std::vector<SparseMatrix *> &)>;
 
-/// @brief Type alias for a function that assembles the HypreParMatrix of a
-/// derivative operator
+/// @brief Type alias for a function that assembles the HypreParMatrix row
+/// blocks of a derivative operator
+///
+/// @see assemble_derivative_sparsematrix_callback_t for the indexing.
 using assemble_derivative_hypreparmatrix_callback_t =
-   std::function<void(HypreParMatrix *&)>;
+   std::function<void(std::vector<HypreParMatrix *> &)>;
 
-/// @brief Type alias for a function that assembles the diagonal of a derivative
-/// operator into an E-vector
-using assemble_diagonal_callback_t = std::function<void(Vector &)>;
+/// @brief Type alias for a function that assembles the diagonal of one row
+/// block of a derivative operator into an E-vector
+///
+/// The first argument names the output (test) field whose row block is wanted;
+/// only a square block has a diagonal at all.
+using assemble_diagonal_callback_t = std::function<void(int, Vector &)>;
 
 /// @brief Type alias for a function that applies the appropriate restriction to
 /// the solution and parameters
@@ -248,60 +258,73 @@ MakeDerivativeHypreParMatrixAssemble(
    std::vector<assemble_derivative_sparsematrix_callback_t> &sparse_callbacks,
    const IntegratorContext &ctx)
 {
-   return [derivative_idx, &sparse_callbacks, ctx](HypreParMatrix *&A)
+   using blocks_t = std::vector<HypreParMatrix *>;
+   return [derivative_idx, &sparse_callbacks, ctx](blocks_t &A)
    {
-      MFEM_VERIFY(ctx.outfds.size() == 1,
-                  "HypreParMatrix assembly requires a single output field");
-
       const size_t trial_field_idx = FindIdx(derivative_idx, ctx.unionfds);
       MFEM_VERIFY(trial_field_idx != SIZE_MAX,
                   "derivative field not found for HypreParMatrix assembly");
 
-      const auto *test_fes_ptr =
-         std::get_if<const ParFiniteElementSpace *>(&ctx.outfds[0].data);
       const auto *trial_fes_ptr =
          std::get_if<const ParFiniteElementSpace *>(
             &ctx.unionfds[trial_field_idx].data);
-      MFEM_VERIFY(test_fes_ptr && *test_fes_ptr,
-                  "HypreParMatrix assembly requires a ParFiniteElementSpace "
-                  "output field");
       MFEM_VERIFY(trial_fes_ptr && *trial_fes_ptr,
                   "HypreParMatrix assembly requires a ParFiniteElementSpace "
                   "derivative field");
-
-      const ParFiniteElementSpace *test_fes = *test_fes_ptr;
       const ParFiniteElementSpace *trial_fes = *trial_fes_ptr;
-      MFEM_VERIFY(test_fes->GetComm() == trial_fes->GetComm(),
-                  "test and trial spaces must use the same MPI communicator");
 
-      SparseMatrix *spmat = nullptr;
+      // Every integrator fills the row blocks it contributes to; the local
+      // blocks are then RAP'd one by one, each with its own test space.
+      std::vector<SparseMatrix *> spmat(ctx.outfds.size(), nullptr);
       for (const auto &f : sparse_callbacks)
       {
          f(spmat);
       }
 
-      MFEM_VERIFY(spmat != nullptr,
+      bool any = false;
+      for (const auto *m : spmat) { any = any || (m != nullptr); }
+      MFEM_VERIFY(any,
                   "internal error: sparse derivative assembly returned NULL");
-      MFEM_VERIFY(spmat->Finalized(),
-                  "local derivative matrix must be finalized");
 
-      if (test_fes == trial_fes)
+      if (A.size() < ctx.outfds.size())
       {
-         HypreParMatrix dA(test_fes->GetComm(), test_fes->GlobalVSize(),
-                           test_fes->GetDofOffsets(), spmat);
-         A = RAP(&dA, test_fes->Dof_TrueDof_Matrix());
-      }
-      else
-      {
-         HypreParMatrix dA(test_fes->GetComm(), test_fes->GlobalVSize(),
-                           trial_fes->GlobalVSize(),
-                           test_fes->GetDofOffsets(),
-                           trial_fes->GetDofOffsets(), spmat);
-         A = RAP(test_fes->Dof_TrueDof_Matrix(), &dA,
-                 trial_fes->Dof_TrueDof_Matrix());
+         A.resize(ctx.outfds.size(), nullptr);
       }
 
-      delete spmat;
+      for (size_t i = 0; i < spmat.size(); i++)
+      {
+         if (spmat[i] == nullptr) { continue; }
+         MFEM_VERIFY(spmat[i]->Finalized(),
+                     "local derivative matrix must be finalized");
+
+         const auto *test_fes_ptr =
+            std::get_if<const ParFiniteElementSpace *>(&ctx.outfds[i].data);
+         MFEM_VERIFY(test_fes_ptr && *test_fes_ptr,
+                     "HypreParMatrix assembly requires a ParFiniteElementSpace "
+                     "output field");
+         const ParFiniteElementSpace *test_fes = *test_fes_ptr;
+         MFEM_VERIFY(test_fes->GetComm() == trial_fes->GetComm(),
+                     "test and trial spaces must use the same "
+                     "MPI communicator");
+
+         if (test_fes == trial_fes)
+         {
+            HypreParMatrix dA(test_fes->GetComm(), test_fes->GlobalVSize(),
+                              test_fes->GetDofOffsets(), spmat[i]);
+            A[i] = RAP(&dA, test_fes->Dof_TrueDof_Matrix());
+         }
+         else
+         {
+            HypreParMatrix dA(test_fes->GetComm(), test_fes->GlobalVSize(),
+                              trial_fes->GlobalVSize(),
+                              test_fes->GetDofOffsets(),
+                              trial_fes->GetDofOffsets(), spmat[i]);
+            A[i] = RAP(test_fes->Dof_TrueDof_Matrix(), &dA,
+                       trial_fes->Dof_TrueDof_Matrix());
+         }
+
+         delete spmat[i];
+      }
    };
 }
 
@@ -563,17 +586,64 @@ public:
       }
    }
 
+   /// @brief Assemble the row blocks of the derivative operator into
+   /// SparseMatrices.
+   ///
+   /// A q-function may write to several output (test) fields. Differentiating
+   /// w.r.t. one field then gives a block column with one row block per output
+   /// field,
+   ///
+   ///     dR/dU = [ dR_U/dU ; dR_Y/dU ],
+   ///
+   /// all sharing the trial space of the differentiated field. Each row block
+   /// is an ordinary test x trial matrix and is assembled separately; the
+   /// stacked matrix, if wanted, is one HypreParMatrixFromBlocks call away.
+   ///
+   /// @param A Resized to the number of output fields and indexed the same way,
+   /// so A[f] is the row block of output field f. Fields whose rows cannot be
+   /// materialised -- quadrature and parameter spaces, which have no basis to
+   /// contract against -- are left as nullptr. Ownership passes to the caller.
+   void Assemble(std::vector<SparseMatrix *> &A)
+   {
+      MFEM_VERIFY(!assemble_derivative_sparsematrix_callbacks.empty(),
+                  "derivative can't be assembled into a SparseMatrix");
+      EnsureQpCache();
+
+      A.assign(outfds.size(), nullptr);
+      for (const auto &f : assemble_derivative_sparsematrix_callbacks)
+      {
+         f(A);
+      }
+   }
+
    /// @brief Assemble the derivative operator into a SparseMatrix.
    ///
+   /// Convenience overload for the common case of a derivative with a single
+   /// assemblable row block. Use Assemble(std::vector<SparseMatrix *> &) when
+   /// the q-function writes to more than one test field.
+   ///
    /// @param A The SparseMatrix to assemble the derivative operator into. Can
    /// be an uninitialized object.
    void Assemble(SparseMatrix *&A)
    {
-      MFEM_ASSERT(!assemble_derivative_sparsematrix_callbacks.empty(),
-                  "derivative can't be assembled into a SparseMatrix");
+      std::vector<SparseMatrix *> blocks;
+      Assemble(blocks);
+      A = SingleBlock(blocks);
+   }
+
+   /// @brief Assemble the row blocks of the derivative operator into
+   /// HypreParMatrices.
+   ///
+   /// @see Assemble(std::vector<SparseMatrix *> &) for the indexing and the
+   /// block structure.
+   void Assemble(std::vector<HypreParMatrix *> &A)
+   {
+      MFEM_VERIFY(!assemble_derivative_hypreparmatrix_callbacks.empty(),
+                  "derivative can't be assembled into a HypreParMatrix");
       EnsureQpCache();
 
-      for (const auto &f : assemble_derivative_sparsematrix_callbacks)
+      A.assign(outfds.size(), nullptr);
+      for (const auto &f : assemble_derivative_hypreparmatrix_callbacks)
       {
          f(A);
       }
@@ -581,18 +651,17 @@ public:
 
    /// @brief Assemble the derivative operator into a HypreParMatrix.
    ///
+   /// Convenience overload for the common case of a derivative with a single
+   /// assemblable row block. Use Assemble(std::vector<HypreParMatrix *> &) when
+   /// the q-function writes to more than one test field.
+   ///
    /// @param A The HypreParMatrix to assemble the derivative operator into. Can
    /// be an uninitialized object.
    void Assemble(HypreParMatrix *&A)
    {
-      MFEM_ASSERT(!assemble_derivative_hypreparmatrix_callbacks.empty(),
-                  "derivative can't be assembled into a HypreParMatrix");
-      EnsureQpCache();
-
-      for (const auto &f : assemble_derivative_hypreparmatrix_callbacks)
-      {
-         f(A);
-      }
+      std::vector<HypreParMatrix *> blocks;
+      Assemble(blocks);
+      A = SingleBlock(blocks);
    }
 
    /// @brief Assemble the derivative of a functional into a Vector.
@@ -635,35 +704,89 @@ public:
       }
    }
 
-   /// @brief Assemble the diagonal of the derivative operator into a T-vector.
+   /// @brief Assemble the diagonal of one row block of the derivative
+   /// operator into a T-vector.
    ///
-   /// @param diag The vector to receive the diagonal (must be T-dof sized).
-   void AssembleDiagonal(Vector &diag) const override
+   /// The derivative is a block column, one row block per output field, all
+   /// sharing the trial space of the differentiated field. Only a block whose
+   /// test space *is* that trial space is square, and only a square block has a
+   /// diagonal. Squareness alone cannot pick the block when several output
+   /// fields live on the trial space, so the row is named rather than inferred.
+   /// Output field ids may differ from the differentiated field's id while
+   /// sharing its ParFiniteElementSpace, in which case several row blocks are
+   /// square at once. Naming the output field id resolves the ambiguity.
+   ///
+   /// @param out_field_id The output field whose row block is wanted.
+   /// @param diag The vector to receive the diagonal, resized to that field's
+   /// T-dof size.
+   void AssembleDiagonal(size_t out_field_id, Vector &diag) const
    {
-      MFEM_ASSERT(!assemble_diagonal_callbacks.empty(),
+      MFEM_VERIFY(!assemble_diagonal_callbacks.empty(),
                   "derivative can't assemble diagonal");
       EnsureQpCache();
-      MFEM_ASSERT(outfds.size() == 1,
-                  "AssembleDiagonal currently requires a single output field");
+
+      const size_t diagonal_idx = FindIdx(out_field_id, outfds);
+      MFEM_VERIFY(diagonal_idx != SIZE_MAX,
+                  "AssembleDiagonal: field " << out_field_id << " is not an "
+                  "output field of this operator");
 
       const auto *test_pf =
-         std::get_if<const ParFiniteElementSpace *>(&outfds[0].data);
+         std::get_if<const ParFiniteElementSpace *>(&outfds[diagonal_idx].data);
       MFEM_VERIFY(test_pf && *test_pf,
                   "AssembleDiagonal: test field must be a ParFiniteElementSpace");
+      const auto *trial_pf =
+         std::get_if<const ParFiniteElementSpace *>(&direction.data);
+      MFEM_VERIFY(trial_pf && *trial_pf,
+                  "AssembleDiagonal: the differentiated field must be a "
+                  "ParFiniteElementSpace");
+      MFEM_VERIFY(*test_pf == *trial_pf,
+                  "AssembleDiagonal: the requested row block is not square and "
+                  "so has no diagonal; its test space is not the trial space "
+                  "of the differentiated field");
 
       prepare_residual(outfds, out_rcache, daction_e);
       for (auto *v : daction_e) { *v = 0.0; }
 
       for (const auto &f : assemble_diagonal_callbacks)
       {
-         f(*daction_e[0]);
+         f(static_cast<int>(out_field_id), *daction_e[diagonal_idx]);
       }
 
       restriction_transpose(outfds, out_rcache, daction_e, daction_l);
-      prolongation_transpose(outfds[0], *daction_l[0], diag);
+      prolongation_transpose(outfds[diagonal_idx],
+                             *daction_l[diagonal_idx], diag);
+   }
+
+   /// @brief Assemble the diagonal of the derivative operator into a T-vector.
+   ///
+   /// Uses the row block of the differentiated field, i.e. the usual diagonal
+   /// block dR_U/dU for GetDerivative(U). Call
+   /// AssembleDiagonal(size_t, Vector &) for any other row block.
+   ///
+   /// @param diag The vector to receive the diagonal.
+   void AssembleDiagonal(Vector &diag) const override
+   {
+      AssembleDiagonal(direction.id, diag);
    }
 
 private:
+   /// The single non-null row block of @a blocks, for the scalar Assemble
+   /// overloads.
+   template <typename mat_t>
+   static mat_t *SingleBlock(const std::vector<mat_t *> &blocks)
+   {
+      mat_t *single = nullptr;
+      int count = 0;
+      for (auto *m : blocks)
+      {
+         if (m != nullptr) { single = m; count++; }
+      }
+      MFEM_VERIFY(count == 1,
+                  "the derivative has " << count << " assemblable row blocks, "
+                  "not one; assemble it into a std::vector instead");
+      return single;
+   }
+
    /// Derivative action callbacks. Depending on the requested derivatives in
    /// DifferentiableOperator the callbacks represent certain combinations of
    /// actions of derivatives of the forward operator.
@@ -1479,16 +1602,23 @@ void DifferentiableOperator::AddIntegrator(
          constexpr size_t derivative_idx = decltype(derivative_id)::value;
          using callback_outputs_t = std::decay_t<decltype(outputs)>;
 
-         bool disable_assemble = false;
+
+         // NOTE: before disable_assemble was looking for any identity outputs,
+         // But this would silently disable assembly for other valid outputs across
+         // different fields.
+         bool any_assemblable_output = false;
          for_constexpr([&](auto j)
          {
-            using output_fop_t = tuple_element_t<j, callback_outputs_t>;
-            if constexpr (is_identity_fop_v<std::decay_t<output_fop_t>>)
+            using output_fop_t =
+               std::decay_t<tuple_element_t<j, callback_outputs_t>>;
+            if constexpr (!is_identity_fop_v<output_fop_t>)
             {
-               disable_assemble = true;
+               any_assemblable_output = true;
             }
          }, std::make_index_sequence<tuple_size<callback_outputs_t>::value> {});
 
+         const bool disable_assemble = !any_assemblable_output;
+
          // Setup the qp cache for the derivative
          setup_callbacks[callback_key].push_back(
             MakeDerivativeSetupCallback(
@@ -1514,13 +1644,21 @@ void DifferentiableOperator::AddIntegrator(
                backend_t::template MakeDerivativeAssemble<derivative_idx>(
                   callback_ctx, qf, inputs, outputs, callback_qp_cache));
 
-            // Assemble the derivative into a HypreParMatrix
-            assemble_hypreparmatrix_callbacks[callback_key].push_back(
-               MakeDerivativeHypreParMatrixAssemble(
-                  derivative_idx,
-                  assemble_sparsematrix_callbacks[callback_key],
-                  callback_ctx));
+            // Assemble the derivative into a HypreParMatrix. This one runs
+            // every sparse callback registered under the key, so it is
+            // registered once and not once per integrator.
+            if (assemble_hypreparmatrix_callbacks[callback_key].empty())
+            {
+               assemble_hypreparmatrix_callbacks[callback_key].push_back(
+                  MakeDerivativeHypreParMatrixAssemble(
+                     derivative_idx,
+                     assemble_sparsematrix_callbacks[callback_key],
+                     callback_ctx));
+            }
+         }
 
+         if (!disable_assemble)
+         {
             // Assemble the diagonal of the derivative into an L-vector
             assemble_diagonal_cbs[callback_key].push_back(
                backend_t::template MakeDerivativeAssembleDiagonal<derivative_idx>(

tests/unit/dfem/test_multiple_outputs.cpp

@@ -174,6 +174,56 @@ struct mass_diffusion_local_qf
    }
 };
 
+// Three outputs across two test fields: V carries mass + diffusion (two
+// outputs on one field), P carries a diffusion scaled by kappa. The
+// two blocks are different, so we should spot if the row blocks
+// get mismatched/corrupted inadvertedly
+constexpr real_t kappa = 2.0;
+
+struct two_field_local_qf
+{
+   inline MFEM_HOST_DEVICE
+   void operator()(
+      const dscalar_t &u,
+      const tensor<dscalar_t, DIM> &dudxi,
+      const tensor<real_t, DIM, DIM> &J,
+      const real_t &w,
+      dscalar_t &out_v,
+      tensor<dscalar_t, DIM> &out_dv,
+      tensor<dscalar_t, DIM> &out_dp) const
+   {
+      const auto invJ = inv(J);
+      const auto detJ = det(J);
+      const auto grad = (dudxi * invJ) * transpose(invJ);
+      out_v = u * detJ * w;
+      out_dv = grad * (detJ * w);
+      out_dp = grad * (kappa * detJ * w);
+   }
+};
+
+// Mass + diffusion on an FE test field, plus quadrature point data on a
+// VectorQuadratureSpace. The second row block has no basis to contract against,
+// so it is not assemblable while the first one still is.
+struct mass_diffusion_qdata_local_qf
+{
+   inline MFEM_HOST_DEVICE
+   void operator()(
+      const dscalar_t &u,
+      const tensor<dscalar_t, DIM> &dudxi,
+      const tensor<real_t, DIM, DIM> &J,
+      const real_t &w,
+      dscalar_t &out1,
+      tensor<dscalar_t, DIM> &out2,
+      tensor<real_t, DIM, DIM> &out3) const
+   {
+      const auto invJ = inv(J);
+      const auto detJ = det(J);
+      out1 = u * detJ * w;
+      out2 = (dudxi * invJ) * transpose(invJ) * (detJ * w);
+      out3 = J;
+   }
+};
+
 TEST_CASE("dFEM Multiple Outputs", "[Parallel][dFEM][GPU]")
 {
    const bool all_tests = launch_all_non_regression_tests;
@@ -496,7 +546,7 @@ TEST_CASE("dFEM Multiple Outputs", "[Parallel][dFEM][GPU]")
          ParBilinearForm blf_fa(&fes);
          blf_fa.AddDomainIntegrator(new MassIntegrator(ir));
          blf_fa.AddDomainIntegrator(new DiffusionIntegrator(ir));
-         blf_fa.SetAssemblyLevel(AssemblyLevel::LEGACYFULL);
+         blf_fa.SetAssemblyLevel(AssemblyLevel::LEGACY);
          blf_fa.Assemble();
          blf_fa.Finalize();
 
@@ -506,6 +556,9 @@ TEST_CASE("dFEM Multiple Outputs", "[Parallel][dFEM][GPU]")
             {COORDINATES, nodes->ParFESpace()},
          };
 
+         // The test field is named V even though it is the same space as the
+         // trial field U, which is the usual "V is the test function" idiom.
+         // The diagonal of dR_V/dU is then named explicitly.
          const std::vector<FieldDescriptor> out_fds
          {
             {V, &fes},
@@ -542,7 +595,7 @@ TEST_CASE("dFEM Multiple Outputs", "[Parallel][dFEM][GPU]")
          SECTION("Multiple Outputs Assemble Diagonal")
          {
             Vector diag(fes.GetTrueVSize()), diag_ref_l(fes.GetVSize());
-            dRdU->AssembleDiagonal(diag);
+            dRdU->AssembleDiagonal(V, diag);
             blf_fa.SpMat().GetDiag(diag_ref_l);
 
             Vector diag_ref(fes.GetTrueVSize());
@@ -556,6 +609,312 @@ TEST_CASE("dFEM Multiple Outputs", "[Parallel][dFEM][GPU]")
             MPI_Barrier(MPI_COMM_WORLD);
          }
       }
+
+      // Outputs spanning two test fields. dR/dU is a block column, one row
+      // block per output field, all sharing the trial space of U, and each row
+      // block assembles into its own matrix.
+      // This would look smth like:
+      //
+      // dR/dU = [ dR_P/dU; dR_U/dU ]
+      //
+      {
+         static constexpr int P = 5;
+
+         ParBilinearForm blf_u(&fes);
+         blf_u.AddDomainIntegrator(new MassIntegrator(ir));
+         blf_u.AddDomainIntegrator(new DiffusionIntegrator(ir));
+         blf_u.SetAssemblyLevel(AssemblyLevel::LEGACY);
+         blf_u.Assemble();
+         blf_u.Finalize();
+
+         ConstantCoefficient kappa_coeff(kappa);
+         ParBilinearForm blf_p(&fes);
+         blf_p.AddDomainIntegrator(new DiffusionIntegrator(kappa_coeff, ir));
+         blf_p.SetAssemblyLevel(AssemblyLevel::LEGACY);
+         blf_p.Assemble();
+         blf_p.Finalize();
+
+         const std::vector<FieldDescriptor> in_fds
+         {
+            {U, &fes},
+            {COORDINATES, nodes->ParFESpace()},
+         };
+
+         // U is deliberately *not* the first output field: A[f] and the
+         // blocks of Mult are indexed by position in out_fds, and
+         // AssembleDiagonal has to pick the block whose field is the
+         // differentiated one rather than simply the first.
+         const std::vector<FieldDescriptor> out_fds
+         {
+            {P, &fes},
+            {U, &fes},
+         };
+
+         DifferentiableOperator dop(in_fds, out_fds, pmesh);
+
+         auto qf = two_field_local_qf{};
+         dop.AddDomainIntegrator<LocalQFBackend>(
+            qf,
+            tuple{Value<U>{}, Gradient<U>{}, Gradient<COORDINATES>{}, Weight{}},
+            tuple{Value<U>{}, Gradient<U>{}, Gradient<P>{}},
+            *ir, all_domain_attr, Derivatives<U> {});
+
+         Vector nodestv;
+         nodes->GetTrueDofs(nodestv);
+         fes.GetRestrictionMatrix()->Mult(x, xtvec);
+         MultiVector X{xtvec, nodestv};
+         auto dRdU = dop.GetDerivative(U, X);
+
+         REQUIRE(dRdU->Height() == 2 * fes.GetTrueVSize());
+         REQUIRE(dRdU->Width() == fes.GetTrueVSize());
+
+         SECTION("Multiple Fields SparseMatrix")
+         {
+            std::vector<SparseMatrix *> A;
+            dRdU->Assemble(A);
+
+            REQUIRE(A.size() == 2);
+            REQUIRE(A[0] != nullptr);
+            REQUIRE(A[1] != nullptr);
+
+            // TestSameMatrices only goes thru the first matrix' sparsity pattern,
+            // so if we compare both ways we can catch entries missing from either side.
+            TestSameMatrices(*A[0], blf_p.SpMat());
+            TestSameMatrices(blf_p.SpMat(), *A[0]);
+            TestSameMatrices(*A[1], blf_u.SpMat());
+            TestSameMatrices(blf_u.SpMat(), *A[1]);
+
+            delete A[0];
+            delete A[1];
+
+            // The element matrix banks are reused, so assembling a second time
+            // (what a Newton loop does every iteration) has to give the same
+            // matrices instead of accumulating on top of the first ones.
+            dRdU->Assemble(A);
+            TestSameMatrices(*A[0], blf_p.SpMat());
+            TestSameMatrices(*A[1], blf_u.SpMat());
+
+            delete A[0];
+            delete A[1];
+            MPI_Barrier(MPI_COMM_WORLD);
+         }
+
+         SECTION("Multiple Fields HypreParMatrix")
+         {
+            std::vector<HypreParMatrix *> A;
+            dRdU->Assemble(A);
+
+            REQUIRE(A.size() == 2);
+            REQUIRE(A[0] != nullptr);
+            REQUIRE(A[1] != nullptr);
+
+            std::unique_ptr<HypreParMatrix> ref_u(blf_u.ParallelAssemble());
+            std::unique_ptr<HypreParMatrix> ref_p(blf_p.ParallelAssemble());
+
+            // The assembled blocks have to agree with the matrix free action
+            // block by block, which is also what pins A[f] to output field f.
+            Vector dir(fes.GetTrueVSize());
+            dir.Randomize(7);
+
+            Vector yp(fes.GetTrueVSize()), yu(fes.GetTrueVSize());
+            MultiVector Y{yp, yu};
+            dRdU->Mult(dir, Y);
+
+            Vector a(fes.GetTrueVSize());
+            A[0]->Mult(dir, a);
+            a -= Y[0];
+            real_t norm_l = a.Normlinf(), norm_g = norm_l;
+            MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX,
+                          pmesh.GetComm());
+            REQUIRE(norm_g == MFEM_Approx(0.0));
+            Vector r(fes.GetTrueVSize());
+            ref_p->Mult(dir, r);
+            a = Y[0];
+            a -= r;
+            norm_l = a.Normlinf();
+            MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX,
+                          pmesh.GetComm());
+            REQUIRE(norm_g == MFEM_Approx(0.0));
+
+            A[1]->Mult(dir, a);
+            a -= Y[1];
+            norm_l = a.Normlinf();
+            MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX,
+                          pmesh.GetComm());
+            REQUIRE(norm_g == MFEM_Approx(0.0));
+            ref_u->Mult(dir, r);
+            a = Y[1];
+            a -= r;
+            norm_l = a.Normlinf();
+            MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX,
+                          pmesh.GetComm());
+            REQUIRE(norm_g == MFEM_Approx(0.0));
+
+            delete A[0];
+            delete A[1];
+            MPI_Barrier(MPI_COMM_WORLD);
+         }
+
+         SECTION("Multiple Fields Assemble Diagonal")
+         {
+            // Both row blocks are square here, since P and U share a space, so
+            // squareness alone cannot pick one. They carry different operators
+            // though, so reading the wrong row block cannot pass unnoticed.
+            const auto check_diag = [&](const Vector &diag,
+                                        ParBilinearForm &ref)
+            {
+               REQUIRE(diag.Size() == fes.GetTrueVSize());
+
+               Vector diag_ref_l(fes.GetVSize());
+               ref.SpMat().GetDiag(diag_ref_l);
+               Vector diag_ref(fes.GetTrueVSize());
+               fes.GetProlongationMatrix()->MultTranspose(diag_ref_l, diag_ref);
+
+               Vector d(diag);
+               d -= diag_ref;
+               real_t norm_l = d.Normlinf(), norm_g = norm_l;
+               MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX,
+                             pmesh.GetComm());
+               REQUIRE(norm_g == MFEM_Approx(0.0));
+            };
+
+            // Default: the row block of the differentiated field, dR_U/dU.
+            Vector diag_u;
+            dRdU->AssembleDiagonal(diag_u);
+            check_diag(diag_u, blf_u);
+
+            // Named: the other row block, dR_P/dU.
+            Vector diag_p;
+            dRdU->AssembleDiagonal(P, diag_p);
+            check_diag(diag_p, blf_p);
+
+            MPI_Barrier(MPI_COMM_WORLD);
+         }
+      }
+
+      // The same two field structure, but now the two output fields live on
+      // different spaces. That is what makes AssembleDiagonal's choice of row
+      // block observable at all: with both fields on one space, reading the
+      // wrong block still lands on an identically sized vector holding the
+      // same numbers, so the check above cannot tell the two apart.
+      {
+         static constexpr int P = 5;
+
+         H1_FECollection fec_p(p + 1, DIM);
+         ParFiniteElementSpace fes_p(&pmesh, &fec_p);
+
+         ParBilinearForm blf_u(&fes);
+         blf_u.AddDomainIntegrator(new MassIntegrator(ir));
+         blf_u.AddDomainIntegrator(new DiffusionIntegrator(ir));
+         blf_u.SetAssemblyLevel(AssemblyLevel::LEGACY);
+         blf_u.Assemble();
+         blf_u.Finalize();
+
+         const std::vector<FieldDescriptor> in_fds
+         {
+            {U, &fes},
+            {COORDINATES, nodes->ParFESpace()},
+         };
+
+         // U is the second output field on purpose, and fes_p is a different
+         // space, so picking out_fds[0] would give a differently sized vector.
+         const std::vector<FieldDescriptor> out_fds
+         {
+            {P, &fes_p},
+            {U, &fes},
+         };
+
+         DifferentiableOperator dop(in_fds, out_fds, pmesh);
+
+         auto qf = two_field_local_qf{};
+         dop.AddDomainIntegrator<LocalQFBackend>(
+            qf,
+            tuple{Value<U>{}, Gradient<U>{}, Gradient<COORDINATES>{}, Weight{}},
+            tuple{Value<U>{}, Gradient<U>{}, Gradient<P>{}},
+            *ir, all_domain_attr, Derivatives<U> {});
+
+         Vector nodestv;
+         nodes->GetTrueDofs(nodestv);
+         fes.GetRestrictionMatrix()->Mult(x, xtvec);
+         MultiVector X{xtvec, nodestv};
+         auto dRdU = dop.GetDerivative(U, X);
+
+         SECTION("Assemble Diagonal Picks The Square Block")
+         {
+            Vector diag;
+            dRdU->AssembleDiagonal(diag);
+
+            // Sized by U's space, not by the first output field's.
+            REQUIRE(diag.Size() == fes.GetTrueVSize());
+
+            Vector diag_ref_l(fes.GetVSize());
+            blf_u.SpMat().GetDiag(diag_ref_l);
+            Vector diag_ref(fes.GetTrueVSize());
+            fes.GetProlongationMatrix()->MultTranspose(diag_ref_l, diag_ref);
+
+            diag -= diag_ref;
+            real_t norm_l = diag.Normlinf(), norm_g = norm_l;
+            MPI_Allreduce(&norm_l, &norm_g, 1, MPI_DOUBLE, MPI_MAX,
+                          pmesh.GetComm());
+            REQUIRE(norm_g == MFEM_Approx(0.0));
+            MPI_Barrier(MPI_COMM_WORLD);
+         }
+      }
+
+      // One assemblable row block and one that is not: S lives on quadrature
+      // points, so it has no basis to contract against and stays null, while
+      // the mass + diffusion block on V assembles as usual.
+      {
+         ParBilinearForm blf_fa(&fes);
+         blf_fa.AddDomainIntegrator(new MassIntegrator(ir));
+         blf_fa.AddDomainIntegrator(new DiffusionIntegrator(ir));
+         blf_fa.SetAssemblyLevel(AssemblyLevel::LEGACYFULL);
+         blf_fa.Assemble();
+         blf_fa.Finalize();
+
+         const std::vector<FieldDescriptor> in_fds
+         {
+            {U, &fes},
+            {COORDINATES, nodes->ParFESpace()},
+         };
+
+         const std::vector<FieldDescriptor> out_fds
+         {
+            {V, &fes},
+            {S, &vqs},
+         };
+
+         DifferentiableOperator dop(in_fds, out_fds, pmesh);
+
+         auto qf = mass_diffusion_qdata_local_qf{};
+         dop.AddDomainIntegrator<LocalQFBackend>(
+            qf,
+            tuple{Value<U>{}, Gradient<U>{}, Gradient<COORDINATES>{}, Weight{}},
+            tuple{Value<V>{}, Gradient<V>{}, Identity<S>{}},
+            *ir, all_domain_attr, Derivatives<U> {});
+
+         Vector nodestv;
+         nodes->GetTrueDofs(nodestv);
+         fes.GetRestrictionMatrix()->Mult(x, xtvec);
+         MultiVector X{xtvec, nodestv};
+         auto dRdU = dop.GetDerivative(U, X);
+
+         SECTION("Partly Assemblable Outputs SparseMatrix")
+         {
+            std::vector<SparseMatrix *> A;
+            dRdU->Assemble(A);
+
+            REQUIRE(A.size() == 2);
+            REQUIRE(A[0] != nullptr);
+            REQUIRE(A[1] == nullptr);
+
+            TestSameMatrices(*A[0], blf_fa.SpMat());
+            TestSameMatrices(blf_fa.SpMat(), *A[0]);
+
+            delete A[0];
+            MPI_Barrier(MPI_COMM_WORLD);
+         }
+      }
    }
 }