diff --git a/README.md b/README.md index e25c4ff..fec922d 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ The package is still in its early stages and many functionalities are still miss - Maps between degrees of freedom and vertices: `vertex_to_dofmap` and `dof_to_vertex` - Blocked Newton Solver - Function evaluation at specified points +- Interpolation matrices from any `ufl.core.expr.Expr` into a compatible space. ## Installation diff --git a/pyproject.toml b/pyproject.toml index cdfb3a7..a7d118d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ build-backend = "scikit_build_core.build" [project] name = "scifem" -version = "0.12.0" +version = "0.13.0" description = "Scientific tools for finite element methods" readme = "README.md" requires-python = ">=3.10" @@ -138,7 +138,7 @@ tag = true sign_tags = false tag_name = "v{new_version}" tag_message = "Bump version: {current_version} → {new_version}" -current_version = "0.12.0" +current_version = "0.13.0" [[tool.bumpversion.files]] diff --git a/src/scifem/__init__.py b/src/scifem/__init__.py index f597a88..6f2cd3a 100644 --- a/src/scifem/__init__.py +++ b/src/scifem/__init__.py @@ -20,6 +20,8 @@ reverse_mark_entities, ) from .eval import evaluate_function +from .interpolation import interpolation_matrix, prepare_interpolation_data + __all__ = [ "PointSource", @@ -43,9 +45,17 @@ "reverse_mark_entities", "norm", "interpolate_function_onto_facet_dofs", + "interpolation_matrix", + "prepare_interpolation_data", ] +if dolfinx.has_petsc4py: + from .interpolation import petsc_interpolation_matrix # type: ignore + + __all__ += ["petsc_interpolation_matrix"] + + def vertex_to_dofmap(V: dolfinx.fem.FunctionSpace) -> npt.NDArray[np.int32]: """ Create a map from the vertices (local to the process) to the correspondning degrees diff --git a/src/scifem/interpolation.py b/src/scifem/interpolation.py new file mode 100644 index 0000000..dec3f34 --- /dev/null +++ b/src/scifem/interpolation.py @@ -0,0 +1,254 @@ +import dolfinx +import ufl +import numpy as np +import numpy.typing as npt +from .utils import unroll_dofmap + +__all__ = ["interpolation_matrix", "prepare_interpolation_data"] + +if dolfinx.has_petsc4py: + from petsc4py import PETSc + + __all__.append("petsc_interpolation_matrix") + + +def prepare_interpolation_data( + expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace +) -> npt.NDArray[np.inexact]: + """Convenience function for preparing data required for assembling the interpolation matrix + + .. math:: + \\begin{align*} + \\Lambda: V &\\rightarrow Q \\\\ + \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j + \\end{align*} + + where :math:`l_j` is the dual basis of the space :math:`Q` with basis functions :math:`\\phi_j`, + and :math:`\\psi_j` are the basis functions of the space :math:`V`. + + Args: + expr: The UFL expression containing a trial function from space `V` + Q: Output interpolation space + Returns: + Interpolation data per cell, as an numpy array. + """ + try: + q_points = Q.element.interpolation_points() + except TypeError: + q_points = Q.element.interpolation_points + + arguments = ufl.algorithms.extract_arguments(expr) + assert len(arguments) == 1 + V = arguments[0].ufl_function_space() + + num_points = q_points.shape[0] + compiled_expr = dolfinx.fem.Expression(expr, q_points) + mesh = Q.mesh + tdim = mesh.topology.dim + num_cells = mesh.topology.index_map(tdim).size_local + # + # (num_cells, num_points, num_dofs*bs, expr_value_size) + array_evaluated = compiled_expr.eval(mesh, np.arange(num_cells, dtype=np.int32)) + assert np.prod(Q.value_shape) == np.prod(expr.ufl_shape) + + im = Q.element.basix_element.interpolation_matrix + + # Get data as (num_cells*num_points,1, expr_shape, num_test_basis_functions*test_block_size) + expr_size = int(np.prod(expr.ufl_shape)) + array_evaluated = array_evaluated.reshape( + num_cells * q_points.shape[0], 1, expr_size, V.dofmap.bs * V.dofmap.dof_layout.num_dofs + ) + jacobian = dolfinx.fem.Expression(ufl.Jacobian(mesh), q_points) + detJ = dolfinx.fem.Expression(ufl.JacobianDeterminant(mesh), q_points) + K = dolfinx.fem.Expression(ufl.JacobianInverse(mesh), q_points) + jacs = jacobian.eval(mesh, np.arange(num_cells, dtype=np.int32)).reshape( + num_cells * num_points, mesh.geometry.dim, mesh.topology.dim + ) + detJs = detJ.eval(mesh, np.arange(num_cells, dtype=np.int32)).flatten() + Ks = K.eval(mesh, np.arange(num_cells, dtype=np.int32)).reshape( + num_cells * num_points, mesh.geometry.dim, mesh.topology.dim + ) + + Q_vs = Q.element.basix_element.value_size + new_array = np.zeros( + (num_cells * num_points, Q.dofmap.bs * Q_vs, V.dofmap.bs * V.dofmap.dof_layout.num_dofs), + dtype=np.float64, + ) + for i in range(V.dofmap.bs * V.dofmap.dof_layout.num_dofs): + for q in range(Q.dofmap.bs): + new_array[:, q * Q_vs : (q + 1) * Q_vs, i] = Q.element.basix_element.pull_back( + array_evaluated[:, :, q * Q_vs : (q + 1) * Q_vs, i], jacs, detJs, Ks + ).reshape(num_cells * num_points, Q_vs) + new_array = new_array.reshape( + num_cells, num_points, Q.dofmap.bs * Q_vs, V.dofmap.bs * V.dofmap.dof_layout.num_dofs + ) + interpolated_matrix = np.zeros( + ( + num_cells, + Q.dofmap.dof_layout.num_dofs * Q.dofmap.bs, + V.dofmap.bs * V.dofmap.dof_layout.num_dofs, + ), + dtype=np.float64, + ) + + for c in range(num_cells): + for i in range(V.dofmap.bs * V.dofmap.dof_layout.num_dofs): + tmp_array = np.zeros((int(num_points), Q.dofmap.bs * Q_vs), dtype=np.float64) + for p in range(num_points): + tmp_array[p] = new_array[c, p, :, i] + if Q.dofmap.bs == 1: + interpolated_matrix[c, :, i] = (im @ tmp_array.T.flatten()).flatten() + else: + for q in range(Q.dofmap.bs): + interpolated_matrix[c, q :: Q.dofmap.bs, i] = ( + im @ tmp_array.T[q].flatten() + ).flatten() + # Apply dof transformation to each column (using Piopla maps) + mesh.topology.create_entity_permutations() + if Q.element.needs_dof_transformations: + cell_perm = mesh.topology.get_cell_permutation_info()[:num_cells] + + permuted_matrix = interpolated_matrix.flatten().copy() + Q.element.Tt_inv_apply( + permuted_matrix, cell_perm, V.dofmap.bs * V.dofmap.dof_layout.num_dofs + ) + else: + permuted_matrix = interpolated_matrix.flatten() + return permuted_matrix.reshape(interpolated_matrix.shape) + + +def interpolation_matrix( + expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace +) -> dolfinx.la.MatrixCSR: + """Create the interpolation matrix :math:`\\Lambda` of a + :py:class:`UFL-expression` such that + + .. math:: + \\begin{align*} + \\Lambda: V &\\rightarrow Q \\\\ + \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j + \\end{align*} + + where :math:`l_j` is the dual basis of the space :math:`Q` with + basis functions :math:`\\phi_j`, and :math:`\\psi_j` are the basis functions of the + space :math:`V`. + + Args: + expr: The UFL expression + Q: Output interpolation space + + Returns: + Interpolation matrix as a :py:class:`MatrixCSR`. + """ + + arguments = ufl.algorithms.extract_arguments(expr) + assert len(arguments) == 1 + V = arguments[0].ufl_function_space() + + interpolation_data = prepare_interpolation_data(expr, Q) + + q = ufl.TestFunction(Q) + a = dolfinx.fem.form(ufl.inner(expr, q) * ufl.dx) + + def scatter( + A: dolfinx.la.MatrixCSR, + num_cells: int, + dofs_visited: npt.NDArray[np.int32], + num_rows_local: int, + array_evaluated: npt.NDArray[np.inexact], + dofmap0: npt.NDArray[np.int32], + dofmap1: npt.NDArray[np.int32], + ): + A.data[:] = 0 + for i in range(num_cells): + rows = dofmap0[i, :] + cols = dofmap1[i, :] + A_local = array_evaluated[i].reshape(len(rows), len(cols)) + row_filter = (dofs_visited[rows] == 1) | (rows >= num_rows_local) + A_local[row_filter] = 0 + A.add(A_local.flatten(), rows, cols) + dofs_visited[rows] = 1 + + A = dolfinx.fem.create_matrix(a) # , dolfinx.la.BlockMode.expanded) + + row_dofmap = unroll_dofmap(Q.dofmap.list, Q.dofmap.bs) # (num_cells, num_rows) + col_dofmap = unroll_dofmap(V.dofmap.list, V.dofmap.bs) # (num_cells, num_cols) + + num_cells = Q.mesh.topology.index_map(Q.mesh.topology.dim).size_local + dofs_visited = np.zeros( + (Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts) * Q.dofmap.index_map_bs, + dtype=np.int8, + ) + num_rows_local = Q.dofmap.index_map.size_local * Q.dofmap.bs + scatter(A, num_cells, dofs_visited, num_rows_local, interpolation_data, row_dofmap, col_dofmap) + A.scatter_reverse() + return A + + +if dolfinx.has_petsc4py: + + def petsc_interpolation_matrix( + expr: ufl.core.expr.Expr, Q: dolfinx.fem.FunctionSpace, use_petsc: bool = False + ) -> PETSc.Mat: + """Create the interpolation matrix :math:`\\Lambda` of a + :py:class:`UFL-expression` such that + + .. math:: + \\begin{align*} + \\Lambda: V &\\rightarrow Q \\\\ + \\Lambda u &= \\sum_{i=0}^{N_Q-1}\\sum_{j=0}^{N_V-1} \\phi_i l_i(expr(\\psi_j))u_j + \\end{align*} + + where :math:`l_j` is the dual basis of the space :math:`Q` with basis + functions :math:`\\phi_j`, and :math:`\\psi_j` are the basis functions + of the space :math:`V`. + + Args: + expr: The UFL expression + Q: Output interpolation space + + Returns: + Interpolation matrix as a :py:class:`PETSc.Mat`. + """ + arguments = ufl.algorithms.extract_arguments(expr) + assert len(arguments) == 1 + V = arguments[0].ufl_function_space() + + interpolation_data = prepare_interpolation_data(expr, Q) + + q = ufl.TestFunction(Q) + a = dolfinx.fem.form(ufl.inner(expr, q) * ufl.dx) + A = dolfinx.fem.petsc.create_matrix(a) + + def scatter( + A: PETSc.Mat, + num_cells: int, + dofs_visited: npt.NDArray[np.int32], + num_rows_local: int, + array_evaluated: npt.NDArray[np.inexact], + dofmap0: npt.NDArray[np.int32], + dofmap1: npt.NDArray[np.int32], + ): + A.zeroEntries() + for i in range(num_cells): + rows = dofmap0[i, :] + cols = dofmap1[i, :] + A_local = array_evaluated[i].reshape(len(rows), len(cols)) + row_filter = (dofs_visited[rows] == 1) | (rows >= num_rows_local) + A_local[row_filter] = 0 + A.setValuesLocal(rows, cols, A_local, addv=PETSc.InsertMode.ADD_VALUES) + dofs_visited[rows] = 1 + + row_dofmap = unroll_dofmap(Q.dofmap.list, Q.dofmap.bs) # (num_cells, num_rows) + col_dofmap = unroll_dofmap(V.dofmap.list, V.dofmap.bs) # (num_cells, num_cols) + num_cells = Q.mesh.topology.index_map(Q.mesh.topology.dim).size_local + dofs_visited = np.zeros( + (Q.dofmap.index_map.size_local + Q.dofmap.index_map.num_ghosts) * Q.dofmap.index_map_bs, + dtype=np.int8, + ) + num_rows_local = Q.dofmap.index_map.size_local * Q.dofmap.bs + scatter( + A, num_cells, dofs_visited, num_rows_local, interpolation_data, row_dofmap, col_dofmap + ) + A.assemble() + return A diff --git a/tests/test_interpolation.py b/tests/test_interpolation.py new file mode 100644 index 0000000..271bddc --- /dev/null +++ b/tests/test_interpolation.py @@ -0,0 +1,196 @@ +from mpi4py import MPI +import dolfinx +import scifem.interpolation +import pytest +import ufl +import numpy as np + + +@pytest.mark.parametrize( + "cell_type", + [ + dolfinx.mesh.CellType.triangle, + dolfinx.mesh.CellType.quadrilateral, + dolfinx.mesh.CellType.tetrahedron, + dolfinx.mesh.CellType.hexahedron, + ], +) +@pytest.mark.parametrize("use_petsc", [True, False]) +@pytest.mark.parametrize("degree", [1, 3, 5]) +def test_interpolation_matrix(use_petsc, cell_type, degree): + if use_petsc: + pytest.importorskip("petsc4py") + + tdim = dolfinx.cpp.mesh.cell_dim(cell_type) + if tdim == 2: + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4, cell_type=cell_type) + elif tdim == 3: + mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 2, 2, 2, cell_type=cell_type) + else: + raise ValueError("Unsupported cell type") + + V = dolfinx.fem.functionspace(mesh, ("DG", degree)) + Q = dolfinx.fem.functionspace(mesh, ("Lagrange", degree)) + + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0] ** degree + x[1] if tdim == 2 else x[0] + x[1] + x[2] ** degree) + + q = dolfinx.fem.Function(Q) + expr = ufl.TrialFunction(V) + if use_petsc: + A = scifem.interpolation.petsc_interpolation_matrix(expr, Q) + A.mult(u.x.petsc_vec, q.x.petsc_vec) + A.destroy() + else: + A = scifem.interpolation.interpolation_matrix(expr, Q) + # Built in matrices has to use a special input vector, with additional ghosts. + _x = dolfinx.la.vector(A.index_map(1), A.block_size[1]) + num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs + _x.array[:num_owned_dofs] = u.x.array[:num_owned_dofs] + _x.scatter_forward() + if not hasattr(dolfinx.la.MatrixCSR, "mult"): + pytest.skip("MatrixCSR has no mult method") + A.mult(_x, q.x) + + q.x.scatter_forward() + + q_ref = dolfinx.fem.Function(Q) + q_ref.interpolate(u) + + np.testing.assert_allclose(q.x.array, q_ref.x.array, rtol=1e-12, atol=1e-13) + + +@pytest.mark.skipif( + not hasattr(dolfinx.fem, "discrete_gradient"), + reason="Cannot verify without discrete gradient from DOLFINx", +) +@pytest.mark.parametrize( + "cell_type", + [ + dolfinx.mesh.CellType.triangle, + dolfinx.mesh.CellType.quadrilateral, + dolfinx.mesh.CellType.tetrahedron, + dolfinx.mesh.CellType.hexahedron, + ], +) +@pytest.mark.parametrize("use_petsc", [True, False]) +@pytest.mark.parametrize("degree", [1, 3, 5]) +def test_discrete_gradient(degree, use_petsc, cell_type): + if use_petsc: + pytest.importorskip("petsc4py") + + tdim = dolfinx.cpp.mesh.cell_dim(cell_type) + if tdim == 2: + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4, cell_type=cell_type) + elif tdim == 3: + mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 2, 2, 2, cell_type=cell_type) + else: + raise ValueError("Unsupported cell type") + + V = dolfinx.fem.functionspace(mesh, ("Lagrange", degree)) + W = dolfinx.fem.functionspace(mesh, ("Nedelec 1st kind H(curl)", degree)) + + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: x[0] ** degree + x[1]) + + w = dolfinx.fem.Function(W) + expr = ufl.grad(ufl.TrialFunction(V)) + + G_ref = dolfinx.fem.discrete_gradient(V, W) + + # Built in matrices has to use a special input vector, with additional ghosts. + try: + _x = dolfinx.la.vector(G_ref.index_map(1), G_ref.block_size[1]) + except AttributeError: + # Bug in DOLFINx 0.9.0 + _x = dolfinx.la.vector(G_ref.index_map(1), G_ref.bs[1]) + + num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs + _x.array[:num_owned_dofs] = u.x.array[:num_owned_dofs] + _x.scatter_forward() + + if use_petsc: + A = scifem.interpolation.petsc_interpolation_matrix(expr, W) + A.mult(u.x.petsc_vec, w.x.petsc_vec) + A.destroy() + else: + if not hasattr(dolfinx.la.MatrixCSR, "mult"): + pytest.skip("MatrixCSR has no mult method") + A = scifem.interpolation.interpolation_matrix(expr, W) + A.mult(_x, w.x) + w.x.scatter_forward() + + w_ref = dolfinx.fem.Function(W) + if not hasattr(dolfinx.la.MatrixCSR, "mult"): + # Fallback to PETSc discrete gradient on 0.9 + pytest.mark.skipif(not dolfinx.has_petsc4py, reason="Cannot verify without petsc4py") + import dolfinx.fem.petsc as _petsc + + G_ref = _petsc.discrete_gradient(V, W) + G_ref.assemble() + G_ref.mult(u.x.petsc_vec, w_ref.x.petsc_vec) + else: + G_ref.mult(_x, w_ref.x) + w_ref.x.scatter_forward() + + np.testing.assert_allclose(w.x.array, w_ref.x.array, rtol=1e-11, atol=1e-12) + + +@pytest.mark.skipif( + not hasattr(dolfinx.fem, "discrete_curl"), + reason="Cannot verify without discrete curl from DOLFINx", +) +@pytest.mark.parametrize( + "cell_type", + [ + dolfinx.mesh.CellType.tetrahedron, + dolfinx.mesh.CellType.hexahedron, + ], +) +@pytest.mark.parametrize("use_petsc", [True, False]) +@pytest.mark.parametrize("degree", [1, 2]) +def test_discrete_curl(degree, use_petsc, cell_type): + if use_petsc: + pytest.importorskip("petsc4py") + + tdim = dolfinx.cpp.mesh.cell_dim(cell_type) + if tdim == 2: + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4, cell_type=cell_type) + elif tdim == 3: + mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 2, 2, 2, cell_type=cell_type) + else: + raise ValueError("Unsupported cell type") + + V = dolfinx.fem.functionspace(mesh, ("Nedelec 2nd kind H(curl)", degree + 1)) + W = dolfinx.fem.functionspace(mesh, ("RT", degree)) + + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: (x[0] ** degree, x[1] ** degree - 1, -x[2])) + + w = dolfinx.fem.Function(W) + expr = ufl.curl(ufl.TrialFunction(V)) + + G_ref = dolfinx.fem.discrete_curl(V, W) + + # Built in matrices has to use a special input vector, with additional ghosts. + _x = dolfinx.la.vector(G_ref.index_map(1), G_ref.block_size[1]) + num_owned_dofs = V.dofmap.index_map.size_local * V.dofmap.index_map_bs + _x.array[:num_owned_dofs] = u.x.array[:num_owned_dofs] + _x.scatter_forward() + + if use_petsc: + A = scifem.interpolation.petsc_interpolation_matrix(expr, W) + A.mult(u.x.petsc_vec, w.x.petsc_vec) + A.destroy() + else: + if not hasattr(dolfinx.la.MatrixCSR, "mult"): + pytest.skip("MatrixCSR has no mult method") + A = scifem.interpolation.interpolation_matrix(expr, W) + A.mult(_x, w.x) + w.x.scatter_forward() + + w_ref = dolfinx.fem.Function(W) + G_ref.mult(_x, w_ref.x) + w_ref.x.scatter_forward() + + np.testing.assert_allclose(w.x.array, w_ref.x.array, rtol=1e-10, atol=1e-11)