Merge branch 'main' into garth/generalise-vector-container

Author: garth-wells

python/dolfinx/fem/function.py

@@ -38,7 +38,9 @@ class Constant(ufl.Constant):
     ]
 
     def __init__(
-        self, domain, c: typing.Union[np.ndarray, Sequence, np.floating, np.complexfloating]
+        self,
+        domain,
+        c: typing.Union[float, np.floating, complex, np.complexfloating, Sequence, np.ndarray],
     ):
         """A constant with respect to a domain.
 

python/dolfinx/wrappers/fem.cpp

@@ -103,7 +103,16 @@ void declare_function_space(nb::module_& m, std::string type)
                      std::shared_ptr<const dolfinx::fem::FiniteElement<T>>>,
                  std::vector<std::shared_ptr<const dolfinx::fem::DofMap>>>(),
              nb::arg("mesh"), nb::arg("elements"), nb::arg("dofmaps"))
-        .def("collapse", &dolfinx::fem::FunctionSpace<T>::collapse)
+        .def("collapse",
+             [](const dolfinx::fem::FunctionSpace<T>& self)
+                 -> std::pair<dolfinx::fem::FunctionSpace<T>,
+                              nanobind::ndarray<std::int32_t, nanobind::numpy>>
+             {
+               auto&& [collapsed_fs, dofs] = self.collapse();
+               return {std::move(collapsed_fs),
+                       dolfinx_wrappers::as_nbarray(std::move(dofs),
+                                                    {dofs.size()})};
+             })
         .def("component", &dolfinx::fem::FunctionSpace<T>::component)
         .def("contains", &dolfinx::fem::FunctionSpace<T>::contains,
              nb::arg("V"))

python/dolfinx/wrappers/mesh.cpp

@@ -696,6 +696,20 @@ void mesh(nb::module_& m)
                 original_cell_indices.data() + original_cell_indices.size());
           },
           nb::arg("original_cell_indices"))
+      .def_prop_ro(
+          "original_cell_indices",
+          [](const dolfinx::mesh::Topology& self)
+          {
+            const std::vector<std::vector<std::int64_t>>& indices
+                = self.original_cell_index;
+            std::vector<nb::ndarray<const std::int64_t, nb::numpy>> idx_nb;
+            for (auto& oci : indices)
+            {
+              idx_nb.push_back(nb::ndarray<const std::int64_t, nb::numpy>(
+                  oci.data(), {oci.size()}));
+            }
+            return idx_nb;
+          })
       .def("connectivity",
            nb::overload_cast<int, int>(&dolfinx::mesh::Topology::connectivity,
                                        nb::const_),

python/test/unit/fem/test_function_space.py

@@ -170,9 +170,13 @@ def test_collapse(W, V):
         Function(W.sub(1))
 
     Ws = [W.sub(i).collapse() for i in range(W.num_sub_spaces)]
-    for Wi, _ in Ws:
+    for Wi, dofs in Ws:
         assert np.allclose(Wi.dofmap.index_map.ghosts, W.dofmap.index_map.ghosts)
 
+        # Number of collapsed dofs in W numbering must agree with the number of dofs
+        # of the collapsed space
+        assert Wi.dofmap.index_map.size_local + Wi.dofmap.index_map.num_ghosts == dofs.size
+
     msh = W.mesh
     cell_imap = msh.topology.index_map(msh.topology.dim)
     num_cells = cell_imap.size_local + cell_imap.num_ghosts