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.