Expose interpolation matrix for MatrixCSR (without PETSc) (#3226)

* Add wrapper to nanobind layer

* Add template parameter

* Add Python frontend

* Add test

python/dolfinx/fem/__init__.py

@@ -12,6 +12,7 @@
 from dolfinx.cpp.fem import create_interpolation_data as _create_interpolation_data
 from dolfinx.cpp.fem import create_sparsity_pattern as _create_sparsity_pattern
 from dolfinx.cpp.fem import discrete_gradient as _discrete_gradient
+from dolfinx.cpp.fem import interpolation_matrix as _interpolation_matrix
 from dolfinx.fem.assemble import (
     apply_lifting,
     assemble_matrix,
@@ -103,6 +104,19 @@ def discrete_gradient(space0: FunctionSpace, space1: FunctionSpace) -> _MatrixCS
     return _discrete_gradient(space0._cpp_object, space1._cpp_object)
 
 
+def interpolation_matrix(space0: FunctionSpace, space1: FunctionSpace) -> _MatrixCSR:
+    """Assemble an interpolation matrix for two function spaces on the same mesh.
+
+    Args:
+        space0: space to interpolate from
+        space1: space to interpolate into
+
+    Returns:
+        Interpolation matrix
+    """
+    return _MatrixCSR(_interpolation_matrix(space0._cpp_object, space1._cpp_object))
+
+
 __all__ = [
     "Constant",
     "Expression",

python/dolfinx/wrappers/assemble.cpp

@@ -43,40 +43,106 @@ namespace nb = nanobind;
 namespace
 {
 
+template <typename U>
+dolfinx::la::SparsityPattern
+create_sparsity(const dolfinx::fem::FunctionSpace<U>& V0,
+                const dolfinx::fem::FunctionSpace<U>& V1)
+{
+  assert(V0.mesh());
+  auto mesh = V0.mesh();
+  assert(V1.mesh());
+  assert(mesh == V1.mesh());
+  MPI_Comm comm = mesh->comm();
+
+  auto dofmap0 = V0.dofmap();
+  assert(dofmap0);
+  auto dofmap1 = V1.dofmap();
+  assert(dofmap1);
+
+  // Create and build  sparsity pattern
+  assert(dofmap0->index_map);
+  assert(dofmap1->index_map);
+  dolfinx::la::SparsityPattern sp(
+      comm, {dofmap1->index_map, dofmap0->index_map},
+      {dofmap1->index_map_bs(), dofmap0->index_map_bs()});
+
+  int tdim = mesh->topology()->dim();
+  auto map = mesh->topology()->index_map(tdim);
+  assert(map);
+  std::vector<std::int32_t> c(map->size_local(), 0);
+  std::iota(c.begin(), c.end(), 0);
+  dolfinx::fem::sparsitybuild::cells(sp, {c, c}, {*dofmap1, *dofmap0});
+  sp.finalize();
+
+  return sp;
+}
+
 // Declare assembler function that have multiple scalar types
 template <typename T, typename U>
 void declare_discrete_operators(nb::module_& m)
 {
+  m.def("interpolation_matrix",
+        [](const dolfinx::fem::FunctionSpace<U>& V0,
+           const dolfinx::fem::FunctionSpace<U>& V1)
+        {
+          // Create sparsity
+          auto sp = create_sparsity(V0, V1);
+
+          // Build operator
+          dolfinx::la::MatrixCSR<T> A(sp);
+
+          auto [bs0, bs1] = A.block_size();
+          if (bs0 == 1 and bs1 == 1)
+          {
+            dolfinx::fem::interpolation_matrix<T, U>(
+                V0, V1, A.template mat_set_values<1, 1>());
+          }
+          else if (bs0 == 2 and bs1 == 1)
+          {
+            dolfinx::fem::interpolation_matrix<T, U>(
+                V0, V1, A.template mat_set_values<2, 1>());
+          }
+          else if (bs0 == 1 and bs1 == 2)
+          {
+            dolfinx::fem::interpolation_matrix<T, U>(
+                V0, V1, A.template mat_set_values<1, 2>());
+          }
+          else if (bs0 == 2 and bs1 == 2)
+          {
+            dolfinx::fem::interpolation_matrix<T, U>(
+                V0, V1, A.template mat_set_values<2, 2>());
+          }
+          else if (bs0 == 3 and bs1 == 1)
+          {
+            dolfinx::fem::interpolation_matrix<T, U>(
+                V0, V1, A.template mat_set_values<3, 1>());
+          }
+          else if (bs0 == 1 and bs1 == 3)
+          {
+            dolfinx::fem::interpolation_matrix<T, U>(
+                V0, V1, A.template mat_set_values<1, 3>());
+          }
+          else if (bs0 == 3 and bs1 == 3)
+          {
+            dolfinx::fem::interpolation_matrix<T, U>(
+                V0, V1, A.template mat_set_values<3, 3>());
+          }
+          else
+          {
+            throw std::runtime_error(
+                "Interpolation matrix not supported between block sizes "
+                + std::to_string(bs0) + " and " + std::to_string(bs1));
+          }
+
+          return A;
+        });
+
   m.def(
       "discrete_gradient",
       [](const dolfinx::fem::FunctionSpace<U>& V0,
          const dolfinx::fem::FunctionSpace<U>& V1)
       {
-        assert(V0.mesh());
-        auto mesh = V0.mesh();
-        assert(V1.mesh());
-        assert(mesh == V1.mesh());
-        MPI_Comm comm = mesh->comm();
-
-        auto dofmap0 = V0.dofmap();
-        assert(dofmap0);
-        auto dofmap1 = V1.dofmap();
-        assert(dofmap1);
-
-        // Create and build  sparsity pattern
-        assert(dofmap0->index_map);
-        assert(dofmap1->index_map);
-        dolfinx::la::SparsityPattern sp(
-            comm, {dofmap1->index_map, dofmap0->index_map},
-            {dofmap1->index_map_bs(), dofmap0->index_map_bs()});
-
-        int tdim = mesh->topology()->dim();
-        auto map = mesh->topology()->index_map(tdim);
-        assert(map);
-        std::vector<std::int32_t> c(map->size_local(), 0);
-        std::iota(c.begin(), c.end(), 0);
-        dolfinx::fem::sparsitybuild::cells(sp, {c, c}, {*dofmap1, *dofmap0});
-        sp.finalize();
+        auto sp = create_sparsity(V0, V1);
 
         // Build operator
         dolfinx::la::MatrixCSR<T> A(sp);

python/test/unit/fem/test_discrete_operators.py

@@ -13,6 +13,7 @@
 
 import dolfinx.la
 import ufl
+from basix.ufl import element
 from dolfinx.fem import Expression, Function, discrete_gradient, functionspace
 from dolfinx.mesh import CellType, GhostMode, create_unit_cube, create_unit_square
 
@@ -188,3 +189,70 @@ def test_gradient_interpolation(cell_type, p, q):
 
     atol = 1000 * np.finfo(dtype).resolution
     assert np.allclose(w_expr.x.array, w.x.array, atol=atol)
+
+
+@pytest.mark.parametrize("p", range(1, 4))
+@pytest.mark.parametrize("q", range(1, 4))
+@pytest.mark.parametrize("from_lagrange", [True, False])
+@pytest.mark.parametrize(
+    "cell_type",
+    [CellType.quadrilateral, CellType.triangle, CellType.tetrahedron, CellType.hexahedron],
+)
+def test_interpolation_matrix(cell_type, p, q, from_lagrange):
+    """Test that discrete interpolation matrix yields the same result as interpolation."""
+    from dolfinx import default_real_type
+    from dolfinx.fem import interpolation_matrix
+
+    comm = MPI.COMM_WORLD
+    if cell_type == CellType.triangle:
+        mesh = create_unit_square(comm, 7, 5, ghost_mode=GhostMode.none, cell_type=cell_type)
+        lagrange = "Lagrange" if from_lagrange else "DG"
+        nedelec = "Nedelec 1st kind H(curl)"
+    elif cell_type == CellType.quadrilateral:
+        mesh = create_unit_square(comm, 11, 6, ghost_mode=GhostMode.none, cell_type=cell_type)
+        lagrange = "Q" if from_lagrange else "DQ"
+        nedelec = "RTCE"
+    elif cell_type == CellType.hexahedron:
+        mesh = create_unit_cube(comm, 3, 2, 1, ghost_mode=GhostMode.none, cell_type=cell_type)
+        lagrange = "Q" if from_lagrange else "DQ"
+        nedelec = "NCE"
+    elif cell_type == CellType.tetrahedron:
+        mesh = create_unit_cube(comm, 3, 2, 2, ghost_mode=GhostMode.none, cell_type=cell_type)
+        lagrange = "Lagrange" if from_lagrange else "DG"
+        nedelec = "Nedelec 1st kind H(curl)"
+    v_el = element(
+        lagrange, mesh.basix_cell(), p, shape=(mesh.geometry.dim,), dtype=default_real_type
+    )
+    s_el = element(nedelec, mesh.basix_cell(), q, dtype=default_real_type)
+    if from_lagrange:
+        el0 = v_el
+        el1 = s_el
+    else:
+        el0 = s_el
+        el1 = v_el
+
+    V = functionspace(mesh, el0)
+    W = functionspace(mesh, el1)
+    G = interpolation_matrix(V, W).to_scipy()
+
+    u = Function(V)
+
+    def f(x):
+        if mesh.geometry.dim == 2:
+            return (x[1] ** p, x[0] ** p)
+        else:
+            return (x[0] ** p, x[2] ** p, x[1] ** p)
+
+    u.interpolate(f)
+    w_vec = Function(W)
+    w_vec.interpolate(u)
+
+    # Compute global matrix vector product
+    w = Function(W)
+    ux = np.zeros(G.shape[1])
+    ux[: len(u.x.array)] = u.x.array[:]
+    w.x.array[: G.shape[0]] = G @ ux
+    w.x.scatter_forward()
+
+    atol = 100 * np.finfo(default_real_type).resolution
+    assert np.allclose(w_vec.x.array, w.x.array, atol=atol)