"""Test of BLACS Redistributor.Requires at least 8 MPI tasks."""import sysimport numpy as npfrom gpaw.band_descriptor import BandDescriptorfrom gpaw.grid_descriptor import GridDescriptorfrom gpaw.mpi import world, distribute_cpusfrom gpaw.utilities.blacs import scalapack_set from gpaw.blacs import BlacsGrid, Redistributor, parallelprint, \ BlacsBandDescriptorG = 120 # number of grid points (G x G x G)N = 10 # number of bands# B: number of band groups# D: number of domainsB = 2D = 2M = N // B # number of bands per groupassert M * B == N, 'M=%d, B=%d, N=%d' % (M,B,N)h = 0.2 # grid spacinga = h * G # side length of box# Set up communicators:domain_comm, kpt_comm, band_comm = distribute_cpus(parsize=D, parsize_bands=B, \ nspins=1, nibzkpts=2)assert world.size >= D*B*kpt_comm.sizeif world.rank == 0: print 'MPI: %d domains, %d band groups, %d kpts' % (domain_comm.size, band_comm.size, kpt_comm.size)# Set up band and grid descriptors:bd = BandDescriptor(N, band_comm, False)gd = GridDescriptor((G, G, G), (a, a, a), True, domain_comm, parsize=D)mcpus, ncpus, blocksize = 2, 2, 6# horrible acronymdef main(seed=42, dtype=float): bbd = BlacsBandDescriptor(world, gd, bd, kpt_comm, mcpus, ncpus, blocksize) nbands = bd.nbands mynbands = bd.mynbands # Note after MPI_Reduce, only meaningful information on gd masters H_Nn = bbd.Nndescriptor.zeros(dtype=dtype) scalapack_set(bbd.Nndescriptor, H_Nn, 0.1, 75.0, 'U') # This is not a BLACS distributed matrix... it should be distinct # on grid masters and then broadcast. C_nN = np.zeros((mynbands, nbands), dtype=dtype) # code below works # H_nn = bbd.nndescriptor.zeros(dtype=dtype) # bbd.Nn2nn.redistribute(H_Nn, H_nn) # parallelprint(world, H_nn) # H_nN = bbd.nNdescriptor.zeros(dtype=dtype) # bbd.nn2nN.redistribute(H_nn, H_nN) # parallelprint(world, H_nN) diagonalizer = bbd.get_diagonalizer() eps_n = np.zeros(bd.mynbands) diagonalizer.diagonalize(H_Nn, C_nN, eps_n) print 'after broadcast' parallelprint(world, C_nN)if __name__ == '__main__': main(dtype=float) # main(dtype=complex)