def test2d_laplacian_periodic(self):
"""
2d array, apply Laplacian, periodic along the two axes
"""
from pnumpy import CubeDecomp
from pnumpy import MultiArrayIter
import operator
from math import sin, pi
# global sizes
ndims = 2
#ns = numpy.array([60] * ndims)
ns = numpy.array([3*4] * ndims)
# local rank and number of procs
rk = MPI.COMM_WORLD.Get_rank()
sz = MPI.COMM_WORLD.Get_size()
# find a domain decomposition
dc = CubeDecomp(sz, ns)
# not all numbers of procs will give a uniform domain decomposition,
# exit if none can be found
if not dc.getDecomp():
if rk == 0:
print('no decomp could be found, adjust the number of procs')
return
# get the local start/stop indices along each axis as a list of
# 1d slices
localSlices = dc.getSlab(rk)
iBeg = numpy.array([s.start for s in localSlices])
iEnd = numpy.array([s.stop for s in localSlices])
nsLocal = numpy.array([s.stop - s.start for s in localSlices])
# create the dist arrays
da = pnumpy.gmdaZeros(nsLocal, numpy.float32, mask=None, numGhosts=1)
laplacian = pnumpy.gmdaZeros(nsLocal, numpy.float32, numGhosts=1)
# set the data
for it in MultiArrayIter(nsLocal):
localInds = it.getIndices()
globalInds = iBeg + localInds
# positions are cell centered, domain is [0, 1]^ndims
position = (globalInds + 0.5)/ numpy.array(ns, numpy.float32)
# sin(2*pi*x) * sin(2*pi*y) ...
da[tuple(localInds)] = reduce(operator.mul,
[numpy.sin(2*numpy.pi*position[i]) for i in range(ndims)])
# apply the Laplacian finite difference operator.
# Start by performing all the operations that do
# not require any communication.
laplacian[:] = 2 * ndims * da
# now subtract the neighbor values which are local to this process
for idim in range(ndims):
# indices shifted in the + direction along axis idim
slabP = [slice(None, None, None) for j in range(idim)] + \
[slice(1, None, None)] + \
[slice(None, None, None) for j in range(idim + 1, ndims)]
# indices shifted in the - direction along axis idim
slabM = [slice(None, None, None) for j in range(idim)] + \
[slice(0, -1, None)] + \
[slice(None, None, None) for j in range(idim + 1, ndims)]
laplacian[slabP] -= da[slabM] # subtract left neighbor
laplacian[slabM] -= da[slabP] # subtract right neighbor
# fetch the data located on other procs
periodic = [True for idim in range(ndims)]
for idim in range(ndims):
# define the positive and negative directions
directionP = tuple([0 for j in range(idim)] + [1] + [0 for j in range(idim + 1, ndims)])
directionM = tuple([0 for j in range(idim)] + [-1] + [0 for j in range(idim + 1, ndims)])
procP = dc.getNeighborProc(rk, directionP, periodic=periodic)
procM = dc.getNeighborProc(rk, directionM, periodic=periodic)
# this is where communication takes place... Note that when
# accessing the data on the low-end side on rank procM we
# access the slide on the positive side on procM (directionP).
# And inversely for the high-end side data...
dataM = da.getData(procM, winID=directionP)
dataP = da.getData(procP, winID=directionM)
# finish off the operator
laplacian[da.getEllipsis(winID=directionM)] -= dataM
laplacian[da.getEllipsis(winID=directionP)] -= dataP
# compute a checksum and send the result to rank 0
checksum = laplacian.reduce(lambda x,y:abs(x) + abs(y), 0.0, rootPe=0)
if rk == 0:
print('checksum = ', checksum)
# float32 calculation has higher error
assert(abs(checksum - 32.0) < 1.e-4)
# free the windows
da.free()
laplacian.free()
#
for disp in (-1, 0), (1, 0), (0, -1), (0, 1):
# negative of disp
nisp = tuple([-d for d in disp])
# local updates
src = domain.shift(disp).getSlice()
dst = domain.shift(nisp).getSlice()
outputData[dst] += inputData[src]
# remote updates
src = domain.extract(disp).getSlice()
dst = domain.extract(nisp).getSlice()
neighRk = dc.getNeighborProc(rk, disp, periodic=notPeriodic)
outputData[dst] += inputData.getData(neighRk, nisp)
# will need to fix the weights when there is no neighbor
if neighRk is None:
numInvalidNeighbors[dst] += 3
#
# south-west, north-west, south-east, north-east
#
for disp in (-1, -1), (-1, 1), (1, -1), (1, 1):
# negative of displacement
nisp = tuple([-d for d in disp])
# compute the star Laplacian in the interior, this does not require
# any communication
laplaceZ = 4 * zda[:]
# local neighbour contributions, no communication
laplaceZ[1:, :] -= zda[0:-1, :]
laplaceZ[0:-1, :] -= zda[1:, :]
laplaceZ[:, 1:] -= zda[:, 0:-1]
laplaceZ[:, 0:-1] -= zda[:, 1:]
# now compute and fill in the halo
# find the procs to the north, east, south, and west. This call will
# return None if there is no neighbour.
noProc = dc.getNeighborProc(rk, (1, 0), periodic=(False, True))
soProc = dc.getNeighborProc(rk, (-1, 0), periodic=(False, True))
eaProc = dc.getNeighborProc(rk, (0, 1), periodic=(False, True))
weProc = dc.getNeighborProc(rk, (0, -1), periodic=(False, True))
# correct at the non-periodic boundaries
if noProc is None:
laplaceZ[-1, :] -= zda[-1, :]
if soProc is None:
laplaceZ[0, :] -= zda[0, :]
# fetch the remote data in the halo of the neighbouring processor. When
# the first argument is None, this amounts to a no-op (zero data are
# returned. Note that winID refers to the neighbour domain. For instance,
# the data to the west of the local domain correspond to the east halo
# on the neighbouring processor.
# any communication
laplaceZ = 4 * zda[:]
# local neighbour contributions, no communication
laplaceZ[1: , :] -= zda[0:-1,:]
laplaceZ[0:-1, :] -= zda[1: ,:]
laplaceZ[:, 1: ] -= zda[:,0:-1]
laplaceZ[:, 0:-1] -= zda[:,1: ]
# now compute and fill in the halo
# find the procs to the north, east, south, and west. This call will
# return None if there is no neighbour.
noProc = dc.getNeighborProc(rk, ( 1, 0), periodic = (False, True))
soProc = dc.getNeighborProc(rk, (-1, 0), periodic = (False, True))
eaProc = dc.getNeighborProc(rk, ( 0, 1), periodic = (False, True))
weProc = dc.getNeighborProc(rk, ( 0, -1), periodic = (False, True))
# correct at the non-periodic boundaries
if noProc is None:
laplaceZ[-1,:] -= zda[-1,:]
if soProc is None:
laplaceZ[0,:] -= zda[0,:]
# fetch the remote data in the halo of the neighbouring processor. When
# the first argument is None, this amounts to a no-op (zero data are
# returned. Note that winID refers to the neighbour domain. For instance,
# the data to the west of the local domain correspond to the east halo
# on the neighbouring processor.
#
for disp in (-1, 0), (1, 0), (0, -1), (0, 1):
# negative of disp
nisp = tuple([-d for d in disp])
# local updates
src = domain.shift(disp).getSlice()
dst = domain.shift(nisp).getSlice()
outputData[dst] += inputData[src]
# remote updates
src = domain.extract(disp).getSlice()
dst = domain.extract(nisp).getSlice()
neighRk = dc.getNeighborProc(rk, disp, periodic=notPeriodic)
outputData[dst] += inputData.getData(neighRk, nisp)
# will need to fix the weights when there is no neighbor
if neighRk is None:
numInvalidNeighbors[dst] += 3
#
# south-west, north-west, south-east, north-east
#
for disp in (-1, -1), (-1, 1), (1, -1), (1, 1):
# negative of displacement
nisp = tuple([-d for d in disp])