Application of point forces, mapping vertex indices to corresponding DOFs

dolfinx
8

Here is a vectorized version of the code above (credit to @dokken)

def vertex_to_dof_map_original(mesh, V):
    num_vertices = (
        mesh.topology.index_map(0).size_local + mesh.topology.index_map(0).num_ghosts
    )
    vertex_to_dof_map = np.empty(num_vertices, dtype=np.int32)
    num_cells = (
        mesh.topology.index_map(mesh.topology.dim).size_local
        + mesh.topology.index_map(mesh.topology.dim).num_ghosts
    )
    c_to_v = mesh.topology.connectivity(mesh.topology.dim, 0)
    vertex_to_dof_map = np.empty(num_vertices, dtype=np.int32)
    num_cells = (
        mesh.topology.index_map(mesh.topology.dim).size_local
        + mesh.topology.index_map(mesh.topology.dim).num_ghosts
    )
    for cell in range(num_cells):
        vertices = c_to_v.links(cell)
        dofs = V.dofmap.cell_dofs(cell)
        for i, vertex in enumerate(vertices):
            vertex_to_dof_map[vertex] = dofs[V.dofmap.dof_layout.entity_dofs(0, i)]
    return vertex_to_dof_map


def vertex_to_dof_map_vectorized(mesh, V):
    num_vertices_per_cell = dolfinx.cpp.mesh.cell_num_entities(
        mesh.topology.cell_type, 0
    )

    dof_layout2 = np.empty((num_vertices_per_cell,), dtype=np.int32)
    for i in range(num_vertices_per_cell):
        var = V.dofmap.dof_layout.entity_dofs(0, i)
        assert len(var) == 1
        dof_layout2[i] = var[0]

    num_vertices = (
        mesh.topology.index_map(0).size_local + mesh.topology.index_map(0).num_ghosts
    )

    c_to_v = mesh.topology.connectivity(mesh.topology.dim, 0)
    assert (
        c_to_v.offsets[1:] - c_to_v.offsets[:-1] == c_to_v.offsets[1]
    ).all(), "Single cell type supported"

    vertex_to_dof_map = np.empty(num_vertices, dtype=np.int32)
    vertex_to_dof_map[c_to_v.array] = V.dofmap.list[:, dof_layout2].reshape(-1)
    return vertex_to_dof_map

The output of the following program

import time
import dolfinx
  
mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 40, 40, 40)

V = dolfinx.fem.functionspace(mesh, ("P", 2))
u = dolfinx.fem.Function(V)


t0 = time.perf_counter()
v2d_original = vertex_to_dof_map_original(mesh, V)
t1 = time.perf_counter()
print("Original: ", t1 - t0)
v2d_vectorized = vertex_to_dof_map_vectorized(mesh, V)
t3 = time.perf_counter()
print("Vectorized: ", t3 - t1)

np.testing.assert_allclose(v2d_original, v2d_vectorized)

is

Original:  4.220731544000955
Vectorized:  0.006541958000525483
1 Like