def main():
# Initialize MPI and find out number of local particles
comm = mpi.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
# number of particles per array
numMyPoints = 1 << 10
numGlobalPoints = size * numMyPoints
avg_vol = 1.0 / numGlobalPoints
dx = numpy.power(avg_vol, 1.0 / dim)
mass = avg_vol
hdx = 1.3
if numGlobalPoints % size != 0:
raise RuntimeError("Run with 2^n num procs!")
# everybody creates two particle arrays with numMyPoints
x1 = random.random(numMyPoints)
y1 = random.random(numMyPoints)
z1 = random.random(numMyPoints)
h1 = numpy.ones_like(x1) * hdx * dx
rho1 = numpy.zeros_like(x1)
x2 = random.random(numMyPoints)
y2 = random.random(numMyPoints)
z2 = random.random(numMyPoints)
h2 = numpy.ones_like(x2) * hdx * dx
rho2 = numpy.zeros_like(x2)
# z1[:] = 1.0
# z2[:] = 0.5
# local particle arrays
pa1 = get_particle_array_wcsph(x=x1, y=y1, h=h1, rho=rho1, z=z1)
pa2 = get_particle_array_wcsph(x=x2, y=y2, h=h2, rho=rho2, z=z2)
# gather the data on root
X1 = numpy.zeros(numGlobalPoints)
Y1 = numpy.zeros(numGlobalPoints)
Z1 = numpy.zeros(numGlobalPoints)
H1 = numpy.ones_like(X1) * hdx * dx
RHO1 = numpy.zeros_like(X1)
gathers = (numpy.ones(size) * numMyPoints, None)
comm.Gatherv(sendbuf=[x1, mpi.DOUBLE], recvbuf=[X1, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[y1, mpi.DOUBLE], recvbuf=[Y1, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[z1, mpi.DOUBLE], recvbuf=[Z1, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[rho1, mpi.DOUBLE],
recvbuf=[RHO1, gathers, mpi.DOUBLE])
X2 = numpy.zeros(numGlobalPoints)
Y2 = numpy.zeros(numGlobalPoints)
Z2 = numpy.zeros(numGlobalPoints)
H2 = numpy.ones_like(X2) * hdx * dx
RHO2 = numpy.zeros_like(X2)
comm.Gatherv(sendbuf=[x2, mpi.DOUBLE], recvbuf=[X2, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[y2, mpi.DOUBLE], recvbuf=[Y2, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[z2, mpi.DOUBLE], recvbuf=[Z2, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[rho2, mpi.DOUBLE],
recvbuf=[RHO2, gathers, mpi.DOUBLE])
# create the particle arrays and PM
PA1 = get_particle_array_wcsph(x=X1, y=Y1, z=Z1, h=H1, rho=RHO1)
PA2 = get_particle_array_wcsph(x=X2, y=Y2, z=Z2, h=H2, rho=RHO2)
# create the parallel manager
PARTICLES = [PA1, PA2]
PM = ZoltanParallelManagerGeometric(dim=dim,
particles=PARTICLES,
comm=comm)
# create the local NNPS object with all the particles
Nnps = BoxSortNNPS(dim=dim, particles=PARTICLES)
Nnps.update()
# only root computes summation density
if rank == 0:
assert numpy.allclose(PA1.rho, 0)
sd_evaluate(Nnps, PM, mass, src_index=1, dst_index=0)
sd_evaluate(Nnps, PM, mass, src_index=0, dst_index=0)
RHO1 = PA1.rho
assert numpy.allclose(PA2.rho, 0)
sd_evaluate(Nnps, PM, mass, src_index=0, dst_index=1)
sd_evaluate(Nnps, PM, mass, src_index=1, dst_index=1)
RHO2 = PA2.rho
# wait for the root...
comm.barrier()
# create the local particle arrays
particles = [pa1, pa2]
# create the local nnps object and parallel manager
pm = ZoltanParallelManagerGeometric(dim=dim,
comm=comm,
particles=particles)
nnps = BoxSortNNPS(dim=dim, particles=particles)
# set the Zoltan parameters (Optional)
pz = pm.pz
pz.set_lb_method("RCB")
pz.Zoltan_Set_Param("DEBUG_LEVEL", "0")
# Update the parallel manager (distribute particles)
pm.update()
# update the local nnps
nnps.update()
# Compute summation density individually on each processor
sd_evaluate(nnps, pm, mass, src_index=0, dst_index=1)
sd_evaluate(nnps, pm, mass, src_index=1, dst_index=1)
sd_evaluate(nnps, pm, mass, src_index=0, dst_index=0)
sd_evaluate(nnps, pm, mass, src_index=1, dst_index=0)
# gather the density and global ids
rho1 = pa1.rho
tmp = comm.gather(rho1)
if rank == 0:
global_rho1 = numpy.concatenate(tmp)
assert (global_rho1.size == numGlobalPoints)
rho2 = pa2.rho
tmp = comm.gather(rho2)
if rank == 0:
global_rho2 = numpy.concatenate(tmp)
assert (global_rho2.size == numGlobalPoints)
# gather global x1 and y1
x1 = pa1.x
tmp = comm.gather(x1)
if rank == 0:
global_x1 = numpy.concatenate(tmp)
assert (global_x1.size == numGlobalPoints)
y1 = pa1.y
tmp = comm.gather(y1)
if rank == 0:
global_y1 = numpy.concatenate(tmp)
assert (global_y1.size == numGlobalPoints)
z1 = pa1.z
tmp = comm.gather(z1)
if rank == 0:
global_z1 = numpy.concatenate(tmp)
assert (global_z1.size == numGlobalPoints)
# gather global x2 and y2
x2 = pa2.x
tmp = comm.gather(x2)
if rank == 0:
global_x2 = numpy.concatenate(tmp)
assert (global_x2.size == numGlobalPoints)
y2 = pa2.y
tmp = comm.gather(y2)
if rank == 0:
global_y2 = numpy.concatenate(tmp)
assert (global_y2.size == numGlobalPoints)
z2 = pa2.z
tmp = comm.gather(z2)
if rank == 0:
global_z2 = numpy.concatenate(tmp)
assert (global_z2.size == numGlobalPoints)
# gather global indices
gid1 = pa1.gid
tmp = comm.gather(gid1)
if rank == 0:
global_gid1 = numpy.concatenate(tmp)
assert (global_gid1.size == numGlobalPoints)
gid2 = pa2.gid
tmp = comm.gather(gid2)
if rank == 0:
global_gid2 = numpy.concatenate(tmp)
assert (global_gid2.size == numGlobalPoints)
# check rho1
if rank == 0:
# make sure the arrays are of the same size
assert (global_x1.size == X1.size)
assert (global_y1.size == Y1.size)
assert (global_z1.size == Z1.size)
for i in range(numGlobalPoints):
# make sure we're chacking the right point
assert abs(global_x1[i] - X1[global_gid1[i]]) < 1e-14
assert abs(global_y1[i] - Y1[global_gid1[i]]) < 1e-14
assert abs(global_z1[i] - Z1[global_gid1[i]]) < 1e-14
diff = abs(global_rho1[i] - RHO1[global_gid1[i]])
condition = diff < 1e-14
assert condition, "diff = %g" % (diff)
# check rho2
if rank == 0:
# make sure the arrays are of the same size
assert (global_x2.size == X2.size)
assert (global_y2.size == Y2.size)
assert (global_z2.size == Z2.size)
for i in range(numGlobalPoints):
# make sure we're chacking the right point
assert abs(global_x2[i] - X2[global_gid2[i]]) < 1e-14
assert abs(global_y2[i] - Y2[global_gid2[i]]) < 1e-14
assert abs(global_z2[i] - Z2[global_gid2[i]]) < 1e-14
diff = abs(global_rho2[i] - RHO2[global_gid2[i]])
condition = diff < 1e-14
assert condition, "diff = %g" % (diff)
if rank == 0:
print("Summation density test: OK")
comm.Gatherv(sendbuf=[x2, mpi.DOUBLE], recvbuf=[X2, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[y2, mpi.DOUBLE], recvbuf=[Y2, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[z2, mpi.DOUBLE], recvbuf=[Z2, gathers, mpi.DOUBLE])
comm.Gatherv(sendbuf=[rho2, mpi.DOUBLE], recvbuf=[RHO2, gathers, mpi.DOUBLE])
# create the particle arrays and PM
PA1 = get_particle_array_wcsph(x=X1, y=Y1, z=Z1, h=H1, rho=RHO1)
PA2 = get_particle_array_wcsph(x=X2, y=Y2, z=Z2, h=H2, rho=RHO2)
# create the parallel manager
PARTICLES = [PA1, PA2]
PM = ZoltanParallelManagerGeometric(dim=dim, particles=PARTICLES, comm=comm)
# create the local NNPS object with all the particles
Nnps = BoxSortNNPS(dim=dim, particles=PARTICLES)
Nnps.update()
# only root computes summation density
if rank == 0:
assert numpy.allclose(PA1.rho, 0)
sd_evaluate(Nnps, PM, mass, src_index=1, dst_index=0)
sd_evaluate(Nnps, PM, mass, src_index=0, dst_index=0)
RHO1 = PA1.rho
assert numpy.allclose(PA2.rho, 0)
sd_evaluate(Nnps, PM, mass, src_index=0, dst_index=1)
sd_evaluate(Nnps, PM, mass, src_index=1, dst_index=1)
RHO2 = PA2.rho
# wait for the root...
comm.barrier()