Add read in partitioning option for adios4dolfinx

[jorgensd]

Just store the adjacencylist generated by forward scatter of cells. This can be directly used as a mesh partitioner.

Ref discussions with @IgorBaratta

[jorgensd]

More detailed description

• use cell index map reverse in scattered to build local partition graph
• store number of processes as attribute in adios2 (first thing to check when reading in mesh in a comm).
• store local range numbers for each set of data (geometry, only owned), (topology, only owned)
• store in local arrays, no big global array, easier to access when reading in, similar to snapshot checkpoint.

[jorgensd]

Proof of concept:

from mpi4py import MPI
import numpy as np
import dolfinx

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, nx=3, ny=3, ghost_mode=dolfinx.mesh.GhostMode.shared_facet)

# Get partitioning
cell_map = mesh.topology.index_map(mesh.topology.dim).index_to_dest_ranks()
num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
cell_offsets = cell_map.offsets[:num_cells_local + 1]
cell_array = cell_map.array[cell_offsets[-1]]

# Compute adjacency with current process as first entry
new_adjacency = np.full(num_cells_local + cell_offsets[-1], -1, dtype=np.int32)
new_offsets = cell_offsets + np.arange(len(cell_offsets), dtype=np.int32)
new_adjacency[new_offsets[:-1]] = MPI.COMM_WORLD.rank
insert_position = np.flatnonzero(new_adjacency == -1)
new_adjacency[insert_position] = cell_array

# Create global connectivity
g_imap = mesh.geometry.index_map()
g_dmap = mesh.geometry.dofmap
num_cells_global = mesh.topology.index_map(mesh.topology.dim).size_global
cell_range = mesh.topology.index_map(mesh.topology.dim).local_range
cmap = mesh.geometry.cmap
geom_layout = cmap.create_dof_layout()
num_dofs_per_cell = geom_layout.num_entity_closure_dofs(mesh.topology.dim)
global_topology = np.zeros((num_cells_local, num_dofs_per_cell), dtype=np.int64)
assert g_dmap.shape[1] == num_dofs_per_cell
global_topology[:, :] = np.asarray(
    g_imap.local_to_global(g_dmap[:num_cells_local, :].reshape(-1))
).reshape(global_topology.shape)

# Extract geometry on process
global_geometry = mesh.geometry.x[:g_imap.size_local].copy()
domain = mesh.ufl_domain()
del mesh

def custom_partitioner(comm: MPI.Intracomm, n, m, topo):
    assert (len(new_offsets) - 1 == topo.num_nodes)

    return dolfinx.graph.adjacencylist(new_adjacency, new_offsets)

new_mesh = dolfinx.mesh.create_mesh(
    MPI.COMM_WORLD, global_topology, global_geometry, domain,
    partitioner=custom_partitioner)
new_cell_map = new_mesh.topology.index_map(new_mesh.topology.dim).index_to_dest_ranks()
np.testing.assert_allclose(cell_map.array, new_cell_map.array)
np.testing.assert_allclose(cell_map.offsets, new_cell_map.offsets)

Should be trivial to add as a boolean option:

1. write_partition_info creates new_adjacency and new_offsets, store number of ranks used to write the mesh, as well as local write ranges of geometry and topology, to be able to access them correctly when reading (store them as local variables, similar to snapshot checkpoints as they should only be possible to read with correct partitioning).
2. read_partition_info, checks file attribute for number of processes used to write function. Reads data using the ranges from local blocks. Read geometry and topology. Create custom partitioner as above. read in mesh with this.

[jorgensd]

Went for a slightly different approach, using the read_dofmap functionality allready implemented. Store an adjacency list (global indices) to file instead of storing local block data.