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.