fenicsx-ii 0.2.0__tar.gz

fenicsx_ii-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Scientific Computing at Simula Research Laboratory
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

fenicsx_ii-0.2.0/PKG-INFO

@@ -0,0 +1,84 @@
+Metadata-Version: 2.4
+Name: fenicsx_ii
+Version: 0.2.0
+Summary: Mixed-dimensional coupling in FEniCSx
+Author-email: "Jørgen S. Dokken" <dokken@simula.no>
+License: MIT License
+        
+        Copyright (c) 2025 Scientific Computing at Simula Research Laboratory
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: fenics-dolfinx>0.9.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: sympy; extra == "test"
+Provides-Extra: dev
+Requires-Dist: pdbpp; extra == "dev"
+Requires-Dist: ipython; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Provides-Extra: docs
+Requires-Dist: jupyter-book; extra == "docs"
+Requires-Dist: jupytext; extra == "docs"
+Requires-Dist: sphinx-codeautolink; extra == "docs"
+Provides-Extra: jupyter
+Requires-Dist: trame_jupyter_extension; extra == "jupyter"
+Requires-Dist: trame; extra == "jupyter"
+Requires-Dist: pyvista[all]>=0.45.0; extra == "jupyter"
+Provides-Extra: all
+Requires-Dist: fenicsx_ii[test]; extra == "all"
+Requires-Dist: fenicsx_ii[dev]; extra == "all"
+Requires-Dist: fenicsx_ii[docs]; extra == "all"
+Requires-Dist: fenicsx_ii[jupyter]; extra == "all"
+Dynamic: license-file
+
+# FEniCSx_ii (FEniCSx trace)
+
+FEniCSx_ii is an extension of FEniCSx that allows users to work with non-conforming 3D-1D meshes.
+
+Given a 3D domain with a function $u\in V(\Omega)$ and a $1D$ domain $\Gamma$.
+We define a restriction operator $\Pi:V\mapsto L^2(\Gamma)$,
+which becomes a central part of the variational formulation.
+
+See for instance
+[D'Angelo & Quarteroni, 2008](https://doi.org/10.1142/S0218202508003108) {cite}`intro-dangelo20083d1d`,
+[Kutcha 2021](https://doi.org/10.1007/978-3-030-55874-1_63) {cite}`intro-kutcha2021trace` or
+[Masri, Kutcha & Riviere, 2024](https://doi.org/10.1137/23M1627390) {cite}`intro-masri2024coupled3d1d`.
+
+The restrictions implemented in {py:mod}`fenicsx_ii` is:
+- {py:class}`PointwiseTrace<fenicsx_ii.PointwiseTrace>`, the operator: $\Pi(u)(\hat x)=u(\hat x)$, $\hat x \in \Gamma$ .
+- {py:class}`Circle<fenicsx_ii.Circle>`, the operator $\Pi(u)(\hat x)=\frac{1}{\vert P_R \vert}\int_{P_{R}(\Gamma(\hat x))}u~\mathrm{d}s$, where $P_R(\Gamma(\hat x))$ is the perimeter of a disk with radius $R$, normal aligning with $\Gamma(\hat x)$ and origin at $\hat x$.
+- {py:class}`Disk<fenicsx_ii.Disk>`, the operator $\Pi(u)(\hat x)=\frac{1}{\vert D_R \vert}\int_{D_R(\Gamma(\hat x))} u~\mathrm{d}x$, where $D_R(\Gamma(\hat x))$ is the disk with radius $R$, normal aligining with $\Gamma(\hat x)$ and origin at $\hat x$.
+
+Any other operator can be implenmented by following the {py:class}`ReductionOperator<fenicsx_ii.ReductionOperator>`-protocol.
+
+## Funding
+The development of FEniCSx_ii has been funded by the Wellcome Trust,
+grant number: 313298/Z/24/Z
+
+## References
+
+```{bibliography}
+:filter: cited
+:labelprefix:
+:keyprefix: intro-
+```

fenicsx_ii-0.2.0/README.md

@@ -0,0 +1,31 @@
+# FEniCSx_ii (FEniCSx trace)
+
+FEniCSx_ii is an extension of FEniCSx that allows users to work with non-conforming 3D-1D meshes.
+
+Given a 3D domain with a function $u\in V(\Omega)$ and a $1D$ domain $\Gamma$.
+We define a restriction operator $\Pi:V\mapsto L^2(\Gamma)$,
+which becomes a central part of the variational formulation.
+
+See for instance
+[D'Angelo & Quarteroni, 2008](https://doi.org/10.1142/S0218202508003108) {cite}`intro-dangelo20083d1d`,
+[Kutcha 2021](https://doi.org/10.1007/978-3-030-55874-1_63) {cite}`intro-kutcha2021trace` or
+[Masri, Kutcha & Riviere, 2024](https://doi.org/10.1137/23M1627390) {cite}`intro-masri2024coupled3d1d`.
+
+The restrictions implemented in {py:mod}`fenicsx_ii` is:
+- {py:class}`PointwiseTrace<fenicsx_ii.PointwiseTrace>`, the operator: $\Pi(u)(\hat x)=u(\hat x)$, $\hat x \in \Gamma$ .
+- {py:class}`Circle<fenicsx_ii.Circle>`, the operator $\Pi(u)(\hat x)=\frac{1}{\vert P_R \vert}\int_{P_{R}(\Gamma(\hat x))}u~\mathrm{d}s$, where $P_R(\Gamma(\hat x))$ is the perimeter of a disk with radius $R$, normal aligning with $\Gamma(\hat x)$ and origin at $\hat x$.
+- {py:class}`Disk<fenicsx_ii.Disk>`, the operator $\Pi(u)(\hat x)=\frac{1}{\vert D_R \vert}\int_{D_R(\Gamma(\hat x))} u~\mathrm{d}x$, where $D_R(\Gamma(\hat x))$ is the disk with radius $R$, normal aligining with $\Gamma(\hat x)$ and origin at $\hat x$.
+
+Any other operator can be implenmented by following the {py:class}`ReductionOperator<fenicsx_ii.ReductionOperator>`-protocol.
+
+## Funding
+The development of FEniCSx_ii has been funded by the Wellcome Trust,
+grant number: 313298/Z/24/Z
+
+## References
+
+```{bibliography}
+:filter: cited
+:labelprefix:
+:keyprefix: intro-
+```

fenicsx_ii-0.2.0/pyproject.toml

@@ -0,0 +1,75 @@
+[build-system] # Require setuptool version due to https://github.com/pypa/setuptools/issues/2938
+requires = ["setuptools>=61.0.0", "wheel"]
+
+[project]
+name = "fenicsx_ii"
+version = "0.2.0"
+description = "Mixed-dimensional coupling in FEniCSx"
+authors = [{ name = "Jørgen S. Dokken", email = "dokken@simula.no" }]
+license = { file = "LICENSE" }
+readme = "README.md"
+dependencies = ["fenics-dolfinx>0.9.0"]
+
+
+[project.optional-dependencies]
+test = ["pytest", "sympy"]
+dev = ["pdbpp", "ipython", "mypy", "ruff"]
+docs = ["jupyter-book", "jupytext", "sphinx-codeautolink"]
+jupyter = ["trame_jupyter_extension", "trame", "pyvista[all]>=0.45.0"]
+all = [
+    "fenicsx_ii[test]",
+    "fenicsx_ii[dev]",
+    "fenicsx_ii[docs]",
+    "fenicsx_ii[jupyter]",
+]
+
+[tool.pytest.ini_options]
+addopts = ["--import-mode=importlib"]
+testpaths = ["tests"]
+
+[tool.mypy]
+ignore_missing_imports = true
+# Folders to exclude
+exclude = ["docs/", "build/"]
+# Folder to check with mypy
+files = ["src", "tests"]
+
+[tool.ruff]
+src = ["src", "docs", "tests", "examples"]
+indent-width = 4
+
+[tool.ruff.lint]
+select = [
+    # Pyflakes
+    "F",
+    # Pycodestyle
+    "E",
+    "W",
+    # isort
+    "I001",
+]
+
+
+[tool.ruff.lint.isort]
+known-first-party = ["fenicsx_iis"]
+known-third-party = [
+    "basix",
+    "dolfinx",
+    "ffcx",
+    "ufl",
+    "gmsh",
+    "numpy",
+    "pytest",
+    "pyadjoint",
+]
+section-order = [
+    "future",
+    "standard-library",
+    "mpi",
+    "third-party",
+    "first-party",
+    "local-folder",
+]
+
+[tool.ruff.lint.isort.sections]
+"mpi" = ["mpi4py", "petsc4py"]

fenicsx_ii-0.2.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

fenicsx_ii-0.2.0/src/fenicsx_ii/__init__.py

@@ -0,0 +1,25 @@
+from .assembly import assemble_scalar, average_coefficients
+from .interpolation import create_interpolation_matrix
+from .matrix_assembler import assemble_matrix, create_matrix
+from .quadrature import Quadrature
+from .restriction_operators import Circle, Disk, PointwiseTrace, ReductionOperator
+from .solver import LinearProblem
+from .ufl_operations import Average
+from .vector_assembler import assemble_vector, create_vector
+
+__all__ = [
+    "Circle",
+    "Disk",
+    "PointwiseTrace",
+    "average_coefficients",
+    "create_interpolation_matrix",
+    "Average",
+    "LinearProblem",
+    "assemble_vector",
+    "create_vector",
+    "assemble_matrix",
+    "create_matrix",
+    "assemble_scalar",
+    "ReductionOperator",
+    "Quadrature",
+]

fenicsx_ii-0.2.0/src/fenicsx_ii/assembly.py

@@ -0,0 +1,103 @@
+import typing
+
+from mpi4py import MPI
+from petsc4py import PETSc
+
+import dolfinx
+import numpy as np
+import ufl
+
+from .interpolation import create_interpolation_matrix
+from .ufl_operations import (
+    AveragedCoefficient,
+    apply_replacer,
+)
+
+__all__ = ["assemble_scalar", "assign_LG_map", "average_coefficients"]
+
+
+def assign_LG_map(
+    C: PETSc.Mat,  # type: ignore
+    row_map: dolfinx.common.IndexMap,
+    col_map: dolfinx.common.IndexMap,
+    row_bs: int,
+    col_bs: int,
+):
+    """
+    Assign a Local to Global map (LG-map) to a PETSc matrix.
+
+    Args:
+        C: The PETSc matrix to which the LG-map is assigned.
+        row_map: The row index map.
+        col_map: The column index map.
+        row_bs: The row block size.
+        col_bs: The column block size.
+    """
+    global_row_map = row_map.local_to_global(
+        np.arange(row_map.size_local + row_map.num_ghosts, dtype=np.int32)
+    ).astype(PETSc.IntType)  # type: ignore
+    global_col_map = col_map.local_to_global(
+        np.arange(col_map.size_local + col_map.num_ghosts, dtype=np.int32)
+    ).astype(PETSc.IntType)  # type: ignore
+    row_map = PETSc.LGMap().create(global_row_map, bsize=row_bs, comm=row_map.comm)  # type: ignore
+    col_map = PETSc.LGMap().create(global_col_map, bsize=col_bs, comm=col_map.comm)  # type: ignore
+    C.setLGMap(row_map, col_map)  # type: ignore
+
+
+def average_coefficients(form: ufl.Form):
+    """Perform the average operation on all coefficients of a given form.
+
+    Note:
+        This updates the underlying vector of each coefficient.
+
+    Args:
+        form: The form whose coefficients are averaged.
+    """
+    for coeff in form.coefficients():
+        if isinstance(coeff, AveragedCoefficient):
+            u = coeff.parent_coefficient
+            res_op = coeff.restriction_operator
+            # Replace rows, i.e. apply interpolation matrix on the left
+            K, _, _ = create_interpolation_matrix(
+                u.function_space,
+                coeff.function_space,
+                res_op,
+                use_petsc=True,
+            )
+            K.mult(u.x.petsc_vec, coeff.x.petsc_vec)
+        coeff.x.scatter_forward()
+
+
+def assemble_scalar(
+    form: ufl.Form,
+    form_compiler_options: dict | None = None,
+    jit_options: dict | None = None,
+    entity_maps: typing.Sequence[dolfinx.mesh.EntityMap] | None = None,
+    op: MPI.Op = MPI.SUM,
+) -> np.inexact | float | complex:
+    """
+    Assemble a rank-0 form into a scalar value. Perform necessary MPI communication.
+
+    Args:
+        form: The rank-0 form to assemble.
+        form_compiler_options: Options for the form compiler.
+        jit_options: Options for just-in-time compilation.
+        entity_maps: Entity maps for the mesh.
+        op: The MPI operation to use for communication (default is
+            :py:obj:`mpi4py.MPI.SUM`).
+    Returns
+        The assembled scalar value.
+    """
+    new_forms = apply_replacer(form)
+    loc_val: float | complex = 0.0
+    for avg_form in new_forms:
+        L_c = dolfinx.fem.form(
+            avg_form,
+            form_compiler_options=form_compiler_options,
+            jit_options=jit_options,
+            entity_maps=entity_maps,
+        )
+        average_coefficients(avg_form)
+        loc_val += dolfinx.fem.assemble_scalar(L_c)
+    val = L_c.mesh.comm.allreduce(loc_val, op=op)
+    return val

fenicsx_ii-0.2.0/src/fenicsx_ii/interpolation.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+from mpi4py import MPI as _MPI
+
+import dolfinx
+import numpy as np
+from dolfinx.common import IndexMap as _im
+
+from .interpolation_utils import create_extended_indexmap, evaluate_basis_function
+from .restriction_operators import ReductionOperator
+from .utils import send_dofs_to_other_process, unroll_dofmap
+
+
+def create_interpolation_matrix(
+    V: dolfinx.fem.FunctionSpace,
+    K: dolfinx.fem.FunctionSpace,
+    red_op: ReductionOperator,
+    tol: float = 1.0e-8,
+    use_petsc: bool = False,
+) -> tuple["PETSc.Mat" | dolfinx.la.MatrixCSR, _im, _im]:  # type: ignore[name-defined] # noqa: F821
+    """
+    Create an interpolation matrix from `V` to `K` with a specific reduction operator
+    applied to the interpolation points of `K`.
+
+    Args:
+        V: The function space to interpolate from.
+        K: The function space to interpolate to.
+        red_op: Reduction operator for each interpolation point on `K`.
+        tol: Tolerance for determining point ownership across processes.
+        use_petsc: Flag to indicate whether to use PETSc for the matrix.
+
+    Returns:
+        The interpolation matrix, as a DOLFINx built in matrix or a PETSc matrix.
+        The accumulation of contributions from all processes has been accounted before
+        returning the matrix.
+    """
+    mesh_to = K.mesh
+    mesh_from = V.mesh
+    num_line_cells = mesh_to.topology.index_map(mesh_to.topology.dim).size_local
+    reference_interpolation_points = K.element.interpolation_points
+
+    num_ip_per_cell = reference_interpolation_points.shape[0]
+    line_cells = np.arange(num_line_cells, dtype=np.int32)
+
+    quadrature_rule = red_op.compute_quadrature(
+        line_cells, reference_interpolation_points
+    )
+    quad_points = quadrature_rule.points
+
+    interpolation_coordinates = quad_points.reshape(-1, mesh_to.geometry.dim)
+    assert mesh_to.geometry.dim == 3
+
+    num_average_qp = red_op.num_points
+
+    point_ownership = dolfinx.geometry.determine_point_ownership(
+        mesh_from, interpolation_coordinates, padding=tol
+    )
+    cells_on_proc = (
+        point_ownership.dest_cells
+    )  # Cells in 3D domain that has interpolation points
+    points_on_proc = (
+        point_ownership.dest_points
+    )  # Interpolation points relating to `cells_on_proc`
+    ip_sender = (
+        point_ownership.src_owner
+    )  # For IP in 1D grid, what process has taken ownership
+    assert (ip_sender >= 0).all()
+    ip_owner = point_ownership.dest_owner  # For received data, who sent it
+
+    num_dofs_per_cell_K = K.dofmap.list.shape[1]
+    incoming_K_dofs, incoming_K_owners = send_dofs_to_other_process(
+        K,
+        ip_owner,
+        ip_sender,
+        np.repeat(np.arange(num_line_cells), num_ip_per_cell * num_average_qp),
+    )
+    # Create extended index map
+    new_imap_K = create_extended_indexmap(
+        K.mesh.comm,
+        K.dofmap.index_map,
+        incoming_K_dofs.flatten(),
+        incoming_K_owners,
+        123,
+    )
+    assert (new_imap_K.global_to_local(incoming_K_dofs.flatten()) >= 0).all()
+
+    incoming_V_dofs, incoming_V_owners = send_dofs_to_other_process(
+        V, ip_sender, ip_owner, cells_on_proc
+    )
+
+    # Create extended index map
+    new_imap_V = create_extended_indexmap(
+        V.mesh.comm,
+        V.dofmap.index_map,
+        incoming_V_dofs.flatten(),
+        incoming_V_owners,
+        321,
+    )
+    num_dofs_per_cell_V = V.dofmap.list.shape[1]
+    new_local_V_dofs = new_imap_V.global_to_local(incoming_V_dofs.flatten())
+    assert (new_local_V_dofs >= 0).all()
+    new_local_V_dofs = new_local_V_dofs.reshape(-1, num_dofs_per_cell_V)
+
+    # Evaluate basis functions in 3D space
+    basis_values_on_V = evaluate_basis_function(V, points_on_proc, cells_on_proc)
+    second_dimension = max(V.dofmap.bs, np.prod(V.element.basix_element.value_shape))
+    recv_basis_functions = np.empty(
+        (len(ip_sender), basis_values_on_V.shape[1], basis_values_on_V.shape[2]),
+        dtype=basis_values_on_V.dtype,
+    )
+    volume_send_to, send_counts_V = np.unique(ip_owner, return_counts=True)
+    line_recv_from, recv_counts_V = np.unique(ip_sender, return_counts=True)
+    basis_send_counts = (
+        send_counts_V * num_dofs_per_cell_V * V.dofmap.bs * second_dimension
+    )
+    basis_recv_counts = (
+        recv_counts_V * num_dofs_per_cell_V * V.dofmap.bs * second_dimension
+    )
+    send_message = [basis_values_on_V.flatten(), basis_send_counts, _MPI.DOUBLE]
+    recv_message = [recv_basis_functions, basis_recv_counts, _MPI.DOUBLE]
+    volume_to_line_comm = mesh_from.comm.Create_dist_graph_adjacent(
+        line_recv_from.tolist(), volume_send_to.tolist(), reorder=False
+    )
+    volume_to_line_comm.Neighbor_alltoallv(send_message, recv_message)
+    # Free communicators post communication
+    volume_to_line_comm.Free()
+
+    # Create sparsity pattern for the interpolation matrix
+    sp = dolfinx.cpp.la.SparsityPattern(
+        K.mesh.comm,
+        [new_imap_K, new_imap_V],
+        [K.dofmap.index_map_bs, V.dofmap.index_map_bs],
+    )
+    insert_position = np.argsort(ip_sender, stable=True)
+    V_in_Q_order = np.argsort(insert_position, stable=True)
+
+    assert K.element.interpolation_ident
+    assert not K.element.needs_dof_transformations
+    for i in range(num_line_cells):
+        local_k_dofs = K.dofmap.list[i]
+        local_v_dofs = new_local_V_dofs[
+            V_in_Q_order[
+                num_dofs_per_cell_K * num_average_qp * i : num_dofs_per_cell_K
+                * num_average_qp
+                * (i + 1)
+            ]
+        ]
+        for j in range(num_dofs_per_cell_K):
+            for k in range(num_average_qp):
+                ldofs = local_v_dofs[j * num_average_qp + k]
+                sp.insert(local_k_dofs[j : j + 1], ldofs)
+    sp.finalize()
+
+    # Create distributed petsc matrix and insert basis function values
+    weights = quadrature_rule.weights
+    scales = quadrature_rule.scales
+
+    if use_petsc:
+        assert dolfinx.has_petsc, (
+            "DOLFINx has to be installed with PETSc support to use PETSc matrices"
+        )
+        from petsc4py import PETSc
+
+        A = dolfinx.cpp.la.petsc.create_matrix(K.mesh.comm, sp)
+
+        def insert_function(A, rows, columns, values):
+            A.setValuesLocal(rows, columns, values, addv=PETSc.InsertMode.ADD)
+
+        def finalize(A):
+            A.assemble()
+    else:
+        A = dolfinx.la.matrix_csr(
+            sp,
+            block_mode=dolfinx.la.BlockMode.compact,
+            dtype=dolfinx.default_scalar_type,
+        )
+
+        def insert_function(A, rows, columns, values):
+            A.add(values, rows, columns)
+
+        def finalize(A):
+            A.scatter_reverse()
+
+    # Keep track of dofs that are local to process, to ensure that we only insert for
+    # - Once per degree of freedom
+    # - Only on the process that owns the degree of freedom
+    dofs_visited = np.full(
+        (K.dofmap.index_map.size_local + K.dofmap.index_map.num_ghosts)
+        * K.dofmap.index_map_bs,
+        False,
+        dtype=np.bool_,
+    )
+    K_bs = K.dofmap.bs
+    dofs_visited[K.dofmap.index_map.size_local * K.dofmap.index_map_bs :] = True
+    padded_K_dm = unroll_dofmap(K.dofmap.list, K_bs)
+    local_visit = np.full(num_average_qp * K_bs, False, dtype=np.bool_)
+    for i in range(num_line_cells):
+        local_k_dofs = padded_K_dm[i]
+        V_slice = V_in_Q_order[
+            num_average_qp * num_dofs_per_cell_K * i : num_average_qp
+            * num_dofs_per_cell_K
+            * (i + 1)
+        ]
+        local_v_dofs = unroll_dofmap(new_local_V_dofs[V_slice], V.dofmap.index_map_bs)
+        local_v_values = recv_basis_functions[V_slice]
+        for j in range(num_dofs_per_cell_K):
+            local_dofs = local_v_dofs[j * num_average_qp : (j + 1) * num_average_qp]
+            local_values = local_v_values[j * num_average_qp : (j + 1) * num_average_qp]
+            average_weights = weights[i * num_dofs_per_cell_K + j]
+            # Get visited dofs from previous run
+            for b in range(K_bs):
+                local_visit[:] = dofs_visited[local_k_dofs[j * K_bs + b]]
+                lv = (
+                    local_values[:, :, b]
+                    * average_weights[:, None]
+                    / scales[i * num_dofs_per_cell_K + j]
+                )
+                for k in range(num_average_qp):
+                    # We insert for all average nodes, thus local visit
+                    # is only updated next time we pass through the `j` loop
+                    lv[k][:] = 0 if local_visit[b * num_average_qp + k] else lv[k]
+                    insert_function(
+                        A,
+                        local_k_dofs[j * K_bs + b : j * K_bs + b + 1],
+                        local_dofs[k],
+                        lv[k],
+                    )
+                    dofs_visited[local_k_dofs[j * K_bs + b]] = True
+    finalize(A)
+    return A, new_imap_K, new_imap_V

fenicsx_ii-0.2.0/src/fenicsx_ii/interpolation_utils.py

@@ -0,0 +1,96 @@
+from mpi4py import MPI as _MPI
+
+import dolfinx
+import numpy as np
+import numpy.typing as npt
+import ufl
+
+
+def evaluate_basis_function(
+    V: dolfinx.fem.FunctionSpace,
+    points: npt.NDArray[np.inexact],
+    cells: npt.NDArray[np.int32],
+    batch_size: int = 200,
+):
+    """Evaluate basis functions in `V` at a set of points.
+
+    Args:
+        V: The function space
+        points: The points to evaluate at (in physical space)
+        cells: The cells. The ith cell corresponds to the ith point.
+        batch_size: The number of points/cells to evaluate at the time,
+            to find a balance between flooding the expression cache and
+            the memory usage.
+
+    Returns:
+        The evaluated basis functions at the given points.
+    """
+    # Pull owning points back to reference cell
+    mesh = V.mesh
+    mesh_nodes = mesh.geometry.x
+    cmap = mesh.geometry.cmap
+
+    ref_x = np.zeros((len(cells), mesh.topology.dim), dtype=mesh.geometry.x.dtype)
+    for i, (point, cell) in enumerate(zip(points, cells)):
+        geom_dofs = mesh.geometry.dofmap[cell]
+        ref_x[i] = cmap.pull_back(point.reshape(-1, 3), mesh_nodes[geom_dofs])
+
+    # Create expression evaluating a trial function (i.e. just the basis function)
+    u = ufl.TestFunction(V)
+    bs = V.dofmap.bs
+    num_dofs = V.dofmap.dof_layout.num_dofs
+    value_size = np.prod(V.element.basix_element.value_shape)
+    if bs > 1 and value_size > 1:
+        raise ValueError(
+            f"A function space cant have both {value_size=} and {bs=} bigger than 1."
+        )
+    # Get basis values as (num_cells, num_basis_functions*bs, max(bs, value_size))
+    basis_values = np.zeros(
+        (len(cells), num_dofs * bs, max(bs, value_size)),
+        dtype=dolfinx.default_scalar_type,
+    )
+    if len(cells) > 0:
+        # NOTE: Expression lives on only this communicator rank
+        # Expression is evaluated for every point in every cell, which means that we
+        # need to discard values that are not on the diagonal.
+        assert ref_x.shape[0] == len(cells)
+        num_batches = len(cells) // batch_size + 1
+        for b in range(num_batches):
+            x_batch = ref_x[b * batch_size : (b + 1) * batch_size]
+            cell_batch = cells[b * batch_size : (b + 1) * batch_size]
+            expr = dolfinx.fem.Expression(u, x_batch, comm=_MPI.COMM_SELF)
+            all_values = expr.eval(mesh, cell_batch)
+            if bs > 1:
+                basis_values[b * batch_size : (b + 1) * batch_size, :, :] = np.swapaxes(
+                    np.diagonal(all_values, axis1=0, axis2=1), 0, 2
+                )
+            else:
+                basis_values[b * batch_size : (b + 1) * batch_size, :, :] = np.diagonal(
+                    all_values, axis1=0, axis2=1
+                ).T.reshape(-1, basis_values.shape[1], basis_values.shape[2])
+    return basis_values
+
+
+def create_extended_indexmap(
+    comm: _MPI.Intracomm,
+    imap: dolfinx.common.IndexMap,
+    potential_new_dofs: npt.NDArray[np.int64],
+    owners: npt.NDArray[np.int32],
+    tag: int,
+) -> dolfinx.common.IndexMap:
+    """
+    Create an extended index map that includes new dofs and their owners.
+    """
+    local_indices = imap.global_to_local(potential_new_dofs)
+    new_ghosts = local_indices < 0
+    new_owners = owners[new_ghosts]
+    extended_ghosts = np.concatenate([imap.ghosts, potential_new_dofs[new_ghosts]])
+    extended_owners = np.concatenate([imap.owners, new_owners])
+
+    return dolfinx.common.IndexMap(
+        comm,
+        imap.size_local,
+        extended_ghosts,
+        extended_owners,
+        tag,
+    )