Improve mapping data on a Function by coordinates

Consider the following minimal example:

from mpi4py import MPI
import dolfinx

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 1, 1)
shape = (2,)
V = dolfinx.fem.functionspace(mesh, ("P", 2, shape))

x = V.tabulate_dof_coordinates()
print(x.shape, mesh.geometry.x.shape)
u = dolfinx.fem.Function(V)
print(len(u.x.array))

Here we create a 1x1 mesh, with 4 vertices.
We create a second order function space, i.e. there will be dofs in 9 locations in the mesh.
However, as we have created a blocked space with shape (2, ) the array u will have 18 dofs.
This is seen by the print:

(9, 3) (4, 3)
18

Let’s next locate the dofs on the lower boundary. They will be in 3 distinct places, but there will be 6 dofs:

import numpy as np

boundary_facets = dolfinx.mesh.locate_entities_boundary(
    mesh, mesh.topology.dim - 1, lambda x: x[1] < 1e-10
)
boundary_dofs = dolfinx.fem.locate_dofs_topological(
    V, mesh.topology.dim - 1, boundary_facets
)
print(boundary_facets, boundary_dofs)

As we can see by this output, boundary dofs only have three entries, as they are blocked.

[0] [0 1 5]

We can now unroll the degrees of freedom, as shown earlier:


bs = V.dofmap.bs
unrolled_dofs = np.empty(len(boundary_dofs) * bs, dtype=np.int32)

unrolled_coords = np.empty((len(boundary_dofs) * bs, 3), dtype=np.float64)

for i, dof in enumerate(boundary_dofs):
    for b in range(bs):
        unrolled_dofs[i * bs + b] = dof * bs + b
        unrolled_coords[i * bs + b] = x[dof]
print(unrolled_dofs, unrolled_coords)

NOTE: I made some minor mistakes in the previous code that I wrote on my phone, and here is the corrected version that yields:

[ 0  1  2  3 10 11] [[ 2.15422807e-17  2.15422807e-17  0.00000000e+00]
 [ 2.15422807e-17  2.15422807e-17  0.00000000e+00]
 [ 1.00000000e+00 -2.15422807e-17  0.00000000e+00]
 [ 1.00000000e+00 -2.15422807e-17  0.00000000e+00]
 [ 5.00000000e-01  0.00000000e+00  0.00000000e+00]
 [ 5.00000000e-01  0.00000000e+00  0.00000000e+00]]