A general parallel process

You want to extract my_alpha based on the color, rather than the rank, otherwise processors in the same group after split end up having different values of alpha.

Modified code

from dolfin import *
# set MPI communicator
comm = MPI.comm_world
size = comm.Get_size()
rank = comm.Get_rank()
# split comm_world in two groups
group_rank = rank // 2
color = rank % 2
new_comm = comm.Split(color, group_rank)
# new_comm = comm.Split(color=color)
new_rank = new_comm.Get_rank()
new_size = new_comm.Get_size()

def poisson(alpha):
    mesh = UnitSquareMesh(new_comm, 256, 256)
    V = FunctionSpace(mesh, 'CG', 1)
    u, v = TrialFunction(V), TestFunction(V)
    bc = DirichletBC(V, Constant(0), 'on_boundary')
    a = inner(Constant(alpha)*grad(u), grad(v))*dx
    L = inner(Constant(1), v)*dx
    uh = Function(V)
    solve(a == L, uh, bc)
    print(f"new_rank = {new_rank}, color={color}, alpha={alpha} -> |uh|=%g {uh.vector().norm('l2')}")

alphas = [3, 8]
# Get alpha based on the color
my_alpha = alphas[color]
poisson(my_alpha)
# mpirun -np 4 python3 file.py

results in

new_rank = 0, color=1, alpha=8 -> |uh|=%g 1.3203509408791925
new_rank = 1, color=1, alpha=8 -> |uh|=%g 1.3203509408791925
new_rank = 0, color=0, alpha=3 -> |uh|=%g 3.5209358423440125
new_rank = 1, color=0, alpha=3 -> |uh|=%g 3.5209358423440125

For future posts, make sure to include the imports in your snippet.

2 Likes