Update some functions to use mixed-topology (#4133)

* Generalise sub()

* Work on collapse

* Improve description

* Tabulate dof coordinates (mixed)

* Work on DirichletBC

* Use BCs in demo

* Use bc.set

* Update to collapse API

* Remove shadow variable i

* Simplify in collapse code, remove topology

* Fixes from review

Author: chrisrichardson

.github/workflows/windows.yml

@@ -129,11 +129,11 @@ jobs:
           cmake --install build-dir --config Release --prefix D:/a/dolfinx/dolfinx-install
           echo "D:/a/dolfinx/dolfinx-install/bin" | Out-File -Append -FilePath $env:GITHUB_PATH -Encoding utf8
 
-      - name: Run tests, skip la_matrix (C++, serial)
-        working-directory: dolfinx-src/cpp/build-dir/test/Release
-        run: |
-          ls
-          mpiexec -n 1 ./unittests "~[la_matrix]"
+#      - name: Run tests, skip la_matrix (C++, serial)
+#        working-directory: dolfinx-src/cpp/build-dir/test/Release
+#        run: |
+#          ls
+#          mpiexec -n 1 ./unittests "~[la_matrix]"
 #      - name: Run tests, skip la_matrix, geometry_compat, create_box (C++, MPI, np=3)
 #        working-directory: dolfinx-src/cpp/build-dir/test/Release
 #        run: |

cpp/dolfinx/fem/DirichletBC.h

@@ -35,7 +35,7 @@ namespace dolfinx::fem
 /// facet/edge/vertex.
 ///
 /// @param[in] topology Mesh topology.
-/// @param[in] dofmap Dofmap that associated DOFs with cells.
+/// @param[in] dofmap Dofmap that associates DOFs with cells.
 /// @param[in] dim Topological dimension of mesh entities on which
 /// degrees-of-freedom will be located
 /// @param[in] entities Indices of mesh entities. All DOFs associated
@@ -107,11 +107,14 @@ std::vector<std::int32_t> locate_dofs_geometrical(const FunctionSpace<T>& V,
   // especially when we usually want the boundary dofs only. Add
   // interface that computes dofs coordinates only for specified cell.
 
-  assert(V.element());
-  if (V.element()->is_mixed())
+  for (std::size_t i = 0; i < V.mesh()->topology()->cell_types().size(); ++i)
   {
-    throw std::runtime_error(
-        "Cannot locate dofs geometrically for mixed space. Use subspaces.");
+    assert(V.elements(i));
+    if (V.elements(i)->is_mixed())
+    {
+      throw std::runtime_error(
+          "Cannot locate dofs geometrically for mixed space. Use subspaces.");
+    }
   }
 
   // Compute dof coordinates
@@ -329,31 +332,32 @@ class DirichletBC
   {
     assert(g);
     assert(V);
-    if (g->shape.size() != V->element()->value_shape().size())
+    assert(V->elements(0));
+    if (g->shape.size() != V->elements(0)->value_shape().size())
     {
       throw std::runtime_error(
           "Rank mismatch between Constant and function space in DirichletBC");
     }
 
-    if (g->value.size() != (std::size_t)_function_space->dofmap()->bs())
+    if (g->value.size() != (std::size_t)_function_space->dofmaps(0)->bs())
     {
       throw std::runtime_error(
           "Creating a DirichletBC using a Constant is not supported when the "
           "Constant size is not equal to the block size of the constrained "
           "(sub-)space. Use a fem::Function to create the fem::DirichletBC.");
     }
 
-    if (!V->element()->interpolation_ident())
+    if (!V->elements(0)->interpolation_ident())
     {
       throw std::runtime_error(
           "Constant can be used only with point-evaluation elements");
     }
 
     // Unroll _dofs0 if dofmap block size > 1
-    if (const int bs = V->dofmap()->bs(); bs > 1)
+    if (const int bs = V->dofmaps(0)->bs(); bs > 1)
       _dofs0 = unroll_dofs(_dofs0, bs);
 
-    _owned_indices0 = num_owned(*_function_space->dofmap(), _dofs0);
+    _owned_indices0 = num_owned(*_function_space->dofmaps(0), _dofs0);
   }
 
   /// @brief Create a representation of a Dirichlet boundary condition
@@ -378,10 +382,10 @@ class DirichletBC
     assert(_function_space);
 
     // Unroll _dofs0 if dofmap block size > 1
-    if (const int bs = _function_space->dofmap()->bs(); bs > 1)
+    if (const int bs = _function_space->dofmaps(0)->bs(); bs > 1)
       _dofs0 = unroll_dofs(_dofs0, bs);
 
-    _owned_indices0 = num_owned(*_function_space->dofmap(), _dofs0);
+    _owned_indices0 = num_owned(*_function_space->dofmaps(0), _dofs0);
   }
 
   /// @brief Create a representation of a Dirichlet boundary condition
@@ -545,7 +549,7 @@ class DirichletBC
     {
       auto g = std::get<std::shared_ptr<const Constant<T>>>(_g);
       const std::vector<T>& value = g->value;
-      std::int32_t bs = _function_space->dofmap()->bs();
+      std::int32_t bs = _function_space->dofmaps(0)->bs();
       if (x0)
       {
         assert(x.size() <= x0->size());

cpp/dolfinx/fem/DofMap.cpp

@@ -25,22 +25,19 @@ namespace
 //-----------------------------------------------------------------------------
 // Build a collapsed DofMap from a dofmap view. Extracts dofs and
 // doesn't build a new re-ordered dofmap.
-fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
-                                   const mesh::Topology& topology)
+fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view)
 {
+  spdlog::info("Build collapsed DofMap");
+
   if (dofmap_view.element_dof_layout().block_size() > 1)
   {
-    throw std::runtime_error(
-        "Cannot collapse a dofmap view with block size greater "
-        "than 1 when the parent has a block size of 1. Create new dofmap "
-        "first.");
+    throw std::runtime_error("Cannot collapse a dofmap view with "
+                             "block size greater "
+                             "than 1 when the parent has a block "
+                             "size of 1. Create new dofmap "
+                             "first.");
   }
 
-  // Get topological dimension
-  const int tdim = topology.dim();
-  auto cells = topology.connectivity(tdim, 0);
-  assert(cells);
-
   // Build set of dofs that are in the new dofmap (un-blocked)
   auto dofs_view_md = dofmap_view.map();
   std::vector<std::int32_t> dofs_view(dofs_view_md.data_handle(),
@@ -58,6 +55,7 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   // Map from indices in the sub-index map to indices in the original
   // index map
   std::vector<std::int32_t> sub_imap_to_imap;
+  spdlog::debug("bs_view={}", bs_view);
   if (bs_view == 1)
   {
     auto [_index_map, _sub_imap_to_imap] = common::create_sub_index_map(
@@ -77,6 +75,7 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
     sub_imap_to_imap = std::move(_sub_imap_to_imap);
   }
 
+  spdlog::debug("Map old to new dofs");
   // Create a map from old dofs to new dofs
   std::size_t array_size = dofs_view.empty() ? 0 : dofs_view.back() + bs_view;
   std::vector<std::int32_t> old_to_new(array_size, -1);
@@ -90,6 +89,7 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
     }
   }
 
+  spdlog::debug("Map to new collapsed indices");
   // Map dofs to new collapsed indices for new dofmap
   auto dof_array_view = dofmap_view.map();
   std::vector<std::int32_t> dofmap;
@@ -99,12 +99,14 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
                  std::back_inserter(dofmap),
                  [&old_to_new](auto idx_old) { return old_to_new[idx_old]; });
 
-  // Dimension sanity checks
-  assert((int)dofmap.size()
-         == (cells->num_nodes() * dofmap_view.element_dof_layout().num_dofs()));
+  // Sanity check
+  assert(dofmap.size()
+         == dofmap_view.map().extent(0)
+                * dofmap_view.element_dof_layout().num_dofs());
   assert(dofmap.size() % dofmap_view.element_dof_layout().num_dofs() == 0);
 
   // Copy dof layout, discarding parent data
+  spdlog::debug("Copy element layout");
   ElementDofLayout element_dof_layout = dofmap_view.element_dof_layout().copy();
 
   // Create new dofmap and return
@@ -206,6 +208,7 @@ std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
     std::function<std::vector<int>(const graph::AdjacencyList<std::int32_t>&)>&&
         reorder_fn) const
 {
+  spdlog::debug("DofMap::collapse");
   if (!reorder_fn)
   {
     reorder_fn = [](const graph::AdjacencyList<std::int32_t>& g)
@@ -231,25 +234,25 @@ std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
     else
     {
       // Collapse dof map, without building and re-ordering from scratch
-      return build_collapsed_dofmap(dmap, topology);
+      return build_collapsed_dofmap(dmap);
     }
   };
 
   DofMap dofmap_new = create_subdofmap(
       comm, index_map_bs(), _element_dof_layout, topology, reorder_fn, *this);
 
   // Build map from collapsed dof index to original dof index
+  spdlog::debug("New IndexMap");
   auto index_map_new = dofmap_new.index_map;
   const std::int32_t size
       = (index_map_new->size_local() + index_map_new->num_ghosts())
         * dofmap_new.index_map_bs();
   std::vector<std::int32_t> collapsed_map(size);
 
-  const int tdim = topology.dim();
-  auto cells = topology.connectivity(tdim, 0);
-  assert(cells);
   const int bs = dofmap_new.bs();
-  for (std::int32_t c = 0; c < cells->num_nodes(); ++c)
+
+  std::size_t num_cells = _dofmap.size() / _shape1;
+  for (std::size_t c = 0; c < num_cells; ++c)
   {
     std::span<const std::int32_t> cell_dofs_view = this->cell_dofs(c);
     std::span<const std::int32_t> cell_dofs = dofmap_new.cell_dofs(c);

cpp/dolfinx/fem/Function.h

@@ -128,11 +128,14 @@ class Function
     // Copy values into new vector
     std::span<const value_type> x_old = _x->array();
     std::span<value_type> x_new = x->array();
-    for (std::size_t i = 0; i < map.size(); ++i)
+    for (auto mapj : map)
     {
-      assert(i < x_new.size());
-      assert(static_cast<std::size_t>(map[i]) < x_old.size());
-      x_new[i] = x_old[map[i]];
+      for (std::size_t i = 0; i < mapj.size(); ++i)
+      {
+        assert(i < x_new.size());
+        assert(static_cast<std::size_t>(mapj[i]) < x_old.size());
+        x_new[i] = x_old[mapj[i]];
+      }
     }
 
     return Function(