from gpaw import parsize, parsize_bandsfrom gpaw.io.hdf5_highlevel import File, HyperslabSelectionfrom gpaw.mpi import worldimport numpy as np# TAU set-upimport pytaupytau.setNode(world.rank)pytau.setNode(world.rank)tau_0 = pytau.profileTimer("Main")tau_1 = pytau.profileTimer("OpenFile")tau_2 = pytau.profileTimer("SetupArray")tau_3 = pytau.profileTimer("Dataset")tau_4 = pytau.profileTimer("Hyperslab")tau_5 = pytau.profileTimer("Write")pytau.start(tau_0)# A = np.ones((25, 1024, 1024), dtype=float) # use about 200 MBpytau.start(tau_1)# parsize is the domain-decomposition, a tuple of three integers# parsize_bands is the state-parallelization, an integerfile = File('psi.hdf5', mode='w', comm=world.get_c_object())pytau.stop(tau_1)# Global dimensions of the psi array. We make it 2D for simplicity:# nbands x Gpytau.start(tau_2)dtype = floatnbands = 600Gx = 20Gy = 20Gz = 20# Check to make sure parallelization options exist and # make sense. Additionally, paralleziation options must# be integer divisible with the dimensions of the psi## For a 64-node partion in VN mode, just hard code# parsize = [4, 4, 4]# parsize_bands = 4# same as partition dimensions#assert parsize is not Noneassert parsize_bands is not Nonedomain_size = parsize[0]*parsize[1]*parsize[2]assert parsize_bands*domain_size == world.size assert nbands % parsize_bands == 0mynbands = nbands / parsize_bandsassert Gx % parsize[0] == 0assert Gy % parsize[1] == 0assert Gz % parsize[2] == 0myGx = Gx / parsize[0]myGy = Gy / parsize[1]myGz = Gz / parsize[2]# Create a simple distributed psi array for testing:# - even ranks in world communicator get 2.0# - odd ranks in world communicator get 0.5 global_shape = (nbands, Gx, Gy, Gz)if world.rank % 2 == 0: factor = 2.0 # even rankselse: factor = 0.5 # odd ranksdata = factor*np.ones((mynbands, myGx, myGy, myGz), dtype)pytau.stop(tau_2)pytau.start(tau_3)dset = file.create_dataset('psi', global_shape, dtype)pytau.stop(tau_3)pytau.start(tau_4)# Calculating indices is the hardest part. # Assume a simple 2D block distribution striding# down the G-index first in the world communicatorband_rank = world.rank / domain_sizedomain_rank = world.rank % domain_sizecoord_x = domain_rank / (parsize[1]*parsize[2])coord_y = (domain_rank - coord_x*parsize[1]*parsize[2]) / parsize[2]coord_z = domain_rank - coord_x*parsize[1]*parsize[2] - coord_y*parsize[2]# slice type is start, stop, stepstep = 1band_indices = slice(band_rank*mynbands, band_rank*mynbands + mynbands, step)domain_indices_x = slice(coord_x*myGx, coord_x*myGx + myGx, step)domain_indices_y = slice(coord_y*myGy, coord_y*myGy + myGy, step)domain_indices_z = slice(coord_z*myGz, coord_z*myGz + myGz, step)indices = (band_indices, domain_indices_x, domain_indices_y, domain_indices_z)print indicesmyslice = HyperslabSelection(indices, dset.shape)pytau.stop(tau_4)pytau.start(tau_5)dset.write(data, selection=myslice, collective=True)# world.barrier() # make sure everyone has finished writingpytau.stop(tau_5)dset.close()file.close()del data # memory should go back to near initial levelspytau.stop(tau_0)