diff --git a/python/dolfinx/fem/__init__.py b/python/dolfinx/fem/__init__.py index 33523fa26a8..2684cb333e9 100644 --- a/python/dolfinx/fem/__init__.py +++ b/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", diff --git a/python/dolfinx/wrappers/assemble.cpp b/python/dolfinx/wrappers/assemble.cpp index a0f8f540177..fb4460b49aa 100644 --- a/python/dolfinx/wrappers/assemble.cpp +++ b/python/dolfinx/wrappers/assemble.cpp @@ -43,40 +43,106 @@ namespace nb = nanobind; namespace { +template +dolfinx::la::SparsityPattern +create_sparsity(const dolfinx::fem::FunctionSpace& V0, + const dolfinx::fem::FunctionSpace& 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 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 void declare_discrete_operators(nb::module_& m) { + m.def("interpolation_matrix", + [](const dolfinx::fem::FunctionSpace& V0, + const dolfinx::fem::FunctionSpace& V1) + { + // Create sparsity + auto sp = create_sparsity(V0, V1); + + // Build operator + dolfinx::la::MatrixCSR A(sp); + + auto [bs0, bs1] = A.block_size(); + if (bs0 == 1 and bs1 == 1) + { + dolfinx::fem::interpolation_matrix( + V0, V1, A.template mat_set_values<1, 1>()); + } + else if (bs0 == 2 and bs1 == 1) + { + dolfinx::fem::interpolation_matrix( + V0, V1, A.template mat_set_values<2, 1>()); + } + else if (bs0 == 1 and bs1 == 2) + { + dolfinx::fem::interpolation_matrix( + V0, V1, A.template mat_set_values<1, 2>()); + } + else if (bs0 == 2 and bs1 == 2) + { + dolfinx::fem::interpolation_matrix( + V0, V1, A.template mat_set_values<2, 2>()); + } + else if (bs0 == 3 and bs1 == 1) + { + dolfinx::fem::interpolation_matrix( + V0, V1, A.template mat_set_values<3, 1>()); + } + else if (bs0 == 1 and bs1 == 3) + { + dolfinx::fem::interpolation_matrix( + V0, V1, A.template mat_set_values<1, 3>()); + } + else if (bs0 == 3 and bs1 == 3) + { + dolfinx::fem::interpolation_matrix( + 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& V0, const dolfinx::fem::FunctionSpace& 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 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 A(sp); diff --git a/python/test/unit/fem/test_discrete_operators.py b/python/test/unit/fem/test_discrete_operators.py index 2eaa9f9b649..c3f66a7fc6d 100644 --- a/python/test/unit/fem/test_discrete_operators.py +++ b/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)