for i in range(size):
s1.axes.plot3D( X[i*numPoints:(i+1)*numPoints],
Y[i*numPoints:(i+1)*numPoints],
Z[i*numPoints:(i+1)*numPoints],
c=colors[i], marker='o', linestyle='None',
alpha=0.5)
s1.axes.set_xlabel( 'X' )
s1.axes.set_ylabel( 'Y' )
s1.axes.set_zlabel( 'Z' )
plt.title('Initital Distribution')
plt.savefig( 'initial.pdf' )
# partition the points using PyZoltan
xa = DoubleArray(numPoints); xa.set_data(x)
ya = DoubleArray(numPoints); ya.set_data(y)
za = DoubleArray(numPoints); za.set_data(z)
gida = UIntArray(numPoints); gida.set_data(gid)
# create the geometric partitioner
pz = zoltan.ZoltanGeometricPartitioner(
dim=3, comm=comm, x=xa, y=ya, z=za, gid=gida)
# call the load balancing function
pz.set_lb_method('RIB')
pz.Zoltan_Set_Param('DEBUG_LEVEL', '1')
pz.Zoltan_LB_Balance()
# get the new assignments
my_global_ids = list( gid )
print ""
# read the input file and distribute objects across processors
numMyPoints, myGlobalIds, x, y = read_input_file()
rank = comm.Get_rank()
parts = np.ones(shape=numMyPoints, dtype=np.int32)
for i in range(numMyPoints):
parts[i] = rank
# now do the load balancing
_x = np.asarray(x); _y = np.asarray(y); _gid = np.asarray(myGlobalIds)
_z = np.zeros_like(x)
x = DoubleArray(numMyPoints); x.set_data(_x)
y = DoubleArray(numMyPoints); y.set_data(_y)
z = DoubleArray(numMyPoints); z.set_data(_z)
gid = UIntArray(numMyPoints); gid.set_data(_gid)
pz = zoltan.ZoltanGeometricPartitioner(dim=2, comm=comm, x=x, y=y, z=z, gid=gid)
# set the weights to 0 by default
weights = pz.weights.get_npy_array()
weights[:] = 0
pz.set_lb_method("RCB")
pz.Zoltan_Set_Param("DEBUG_LEVEL","0")
pz.Zoltan_LB_Balance()
if rank == 0:
print()
# read the input file and distribute objects across processors
numMyPoints, myGlobalIds, x, y = read_input_file()
rank = comm.Get_rank()
parts = np.ones(shape=numMyPoints, dtype=np.int32)
for i in range(numMyPoints):
parts[i] = rank
# now do the load balancing
_x = np.asarray(x); _y = np.asarray(y); _gid = np.asarray(myGlobalIds)
_z = np.zeros_like(x)
x = DoubleArray(numMyPoints); x.set_data(_x)
y = DoubleArray(numMyPoints); y.set_data(_y)
z = DoubleArray(numMyPoints); z.set_data(_z)
gid = UIntArray(numMyPoints); gid.set_data(_gid)
pz = zoltan.ZoltanGeometricPartitioner(dim=2, comm=comm, x=x, y=y, z=z, gid=gid)
# set the weights to 0 by default
weights = pz.weights.get_npy_array()
weights[:] = 0
pz.set_lb_method("RCB")
pz.Zoltan_Set_Param("DEBUG_LEVEL","0")
pz.Zoltan_LB_Balance()
if rank == 0:
comm.Gather(sendbuf=x, recvbuf=X, root=0)
comm.Gather(sendbuf=y, recvbuf=Y, root=0)
comm.Gather(sendbuf=z, recvbuf=Z, root=0)
if rank == 0:
slice_data = [
slice(i * numPoints, (i + 1) * numPoints) for i in range(size)
]
plot_points(X,
Y,
Z,
slice_data,
title="Initial Distribution",
filename="initial.pdf")
# partition the points using PyZoltan
xa = DoubleArray(numPoints)
xa.set_data(x)
ya = DoubleArray(numPoints)
ya.set_data(y)
za = DoubleArray(numPoints)
za.set_data(z)
gida = UIntArray(numPoints)
gida.set_data(gid)
# create the geometric partitioner
pz = zoltan.ZoltanGeometricPartitioner(dim=3,
comm=comm,
x=xa,
y=ya,
z=za,
gid=gida)