Partitioned T-joint

minimal_example.py

import dolfinx
import ufl
import basix
import numpy as np
from mpi4py import MPI

comm = MPI.COMM_WORLD

#         (3)
#          |
#          |[2]
#          |
# (1)-----(0)-----(2)
#     [0]     [1]


def compute_local_range(comm: MPI.Intracomm, N: np.int64):
    """
    Divide a set of `N` objects into `M` partitions, where `M` is
    the size of the MPI communicator `comm`.

    NOTE: If N is not divisible by the number of ranks, the first `r`
    processes gets an extra value

    Returns the local range of values
    """
    rank = comm.rank
    size = comm.size
    n = N // size
    r = N % size
    # First r processes has one extra value
    if rank < r:
        return [rank * (n + 1), (rank + 1) * (n + 1)]
    else:
        return [rank * n + r, (rank + 1) * n + r]


coord = np.array([[0.0, 0.0, 0.0], [-1.0, 0.0, 0.0], [+1.0, 0.0, 0.0], [0.0, +1.0, 0.0]])
cells = np.array([[0, 1], [0, 2], [0, 3]])

cell_range = compute_local_range(comm, cells.shape[0])
node_range = compute_local_range(comm, coord.shape[0])


local_cells = cells[cell_range[0]:cell_range[1], :]

ufl_elem = basix.ufl.element("Lagrange", "interval", 1, shape=(3, ))
ufl_mesh = ufl.Mesh(ufl_elem)

# mesh = dolfinx.mesh.create_mesh(comm, cells, coord, ufl_mesh)

# FIXME: Get vertex extraction to work properly here, as now we asssume topology is same as geometry
boundary_vertices = np.unique(local_cells.flatten()).astype(np.int64)
topology = dolfinx.cpp.mesh.create_topology(MPI.COMM_WORLD, [dolfinx.mesh.CellType.interval], [local_cells.flatten()], [np.arange(cells.shape[0], dtype=np.int64)], [np.zeros(0, dtype=np.int64)], boundary_vertices)

c_el = dolfinx.fem.coordinate_element(ufl_elem.basix_element)
nodes = np.unique(local_cells.flatten()).astype(np.int64)
xdofs = local_cells.flatten()
x= coord[nodes].flatten()


# FIXME: locate nodes that are in `local_cells` and not owned by this rank


geometry = dolfinx.cpp.mesh.create_geometry(topology, [c_el._cpp_object], nodes, 
		xdofs,
		x,
	3)

if x.dtype == np.float64:
	cpp_mesh = dolfinx.cpp.mesh.Mesh_float64(comm, topology, geometry)
elif x.dtype == np.float32:
	cpp_mesh = dolfinx.cpp.mesh.Mesh_float32(comm, topology, geometry)
else:
	raise RuntimeError(f"Unsupported dtype for mesh {points.dtype}")
mesh = dolfinx.mesh.Mesh(cpp_mesh, domain = ufl_mesh)

for tdim in [0, 1]:
	imap = mesh.topology.index_map(tdim)
	lidx = np.arange(imap.size_local, dtype=np.int32)
	size = imap.size_global
	gidx = imap.local_to_global(lidx)
	print(f"(tdim={tdim}, rank={comm.rank}): {imap.size_local=}, {imap.size_global=} global_indices = {gidx}")


with dolfinx.io.XDMFFile(comm, "mwe_split.xdmf", "w") as xdmf:
	xdmf.write_mesh(mesh)