def redistribute(self, in_wGG, out_x=None):
"""Redistribute array.
Switch between two kinds of parallel distributions:
1) parallel over G-vectors (second dimension of in_wGG)
2) parallel over frequency (first dimension of in_wGG)
Returns new array using the memory in the 1-d array out_x.
"""
comm = self.blockcomm
if comm.size == 1:
return in_wGG
nw = len(self.omega_w)
nG = in_wGG.shape[2]
mynw = (nw + comm.size - 1) // comm.size
mynG = (nG + comm.size - 1) // comm.size
bg1 = BlacsGrid(comm, comm.size, 1)
bg2 = BlacsGrid(comm, 1, comm.size)
md1 = BlacsDescriptor(bg1, nw, nG**2, mynw, nG**2)
md2 = BlacsDescriptor(bg2, nw, nG**2, nw, mynG * nG)
if len(in_wGG) == nw:
mdin = md2
mdout = md1
else:
mdin = md1
mdout = md2
r = Redistributor(comm, mdin, mdout)
outshape = (mdout.shape[0], mdout.shape[1] // nG, nG)
if out_x is None:
out_wGG = np.empty(outshape, complex)
else:
out_wGG = out_x[:np.product(outshape)].reshape(outshape)
r.redistribute(in_wGG.reshape(mdin.shape),
out_wGG.reshape(mdout.shape))
return out_wGG
def distribute_frequencies(self, chi0_wGG):
"""Distribute frequencies to all cores."""
world = self.world
comm = self.blockcomm
if world.size == 1:
return chi0_wGG
nw = len(self.omega_w)
nG = chi0_wGG.shape[2]
mynw = (nw + world.size - 1) // world.size
mynG = (nG + comm.size - 1) // comm.size
wa = min(world.rank * mynw, nw)
wb = min(wa + mynw, nw)
if self.blockcomm.size == 1:
return chi0_wGG[wa:wb].copy()
if self.kncomm.rank == 0:
bg1 = BlacsGrid(comm, 1, comm.size)
in_wGG = chi0_wGG.reshape((nw, -1))
else:
bg1 = DryRunBlacsGrid(mpi.serial_comm, 1, 1)
in_wGG = np.zeros((0, 0), complex)
md1 = BlacsDescriptor(bg1, nw, nG**2, nw, mynG * nG)
bg2 = BlacsGrid(world, world.size, 1)
md2 = BlacsDescriptor(bg2, nw, nG**2, mynw, nG**2)
r = Redistributor(world, md1, md2)
shape = (wb - wa, nG, nG)
out_wGG = np.empty(shape, complex)
r.redistribute(in_wGG, out_wGG.reshape((wb - wa, nG**2)))
return out_wGG
def diagonalize_full_hamiltonian(self, ham, atoms, occupations, txt,
nbands=None, scalapack=None, expert=False):
assert self.dtype == complex
if nbands is None:
nbands = self.pd.ngmin // self.bd.comm.size * self.bd.comm.size
else:
assert nbands <= self.pd.ngmin
if expert:
iu = nbands
else:
iu = None
self.bd = bd = BandDescriptor(nbands, self.bd.comm)
p = functools.partial(print, file=txt)
p('Diagonalizing full Hamiltonian ({0} lowest bands)'.format(nbands))
p('Matrix size (min, max): {0}, {1}'.format(self.pd.ngmin,
self.pd.ngmax))
mem = 3 * self.pd.ngmax**2 * 16 / bd.comm.size / 1024**2
p('Approximate memory usage per core: {0:.3f} MB'.format(mem))
if bd.comm.size > 1:
if isinstance(scalapack, (list, tuple)):
nprow, npcol, b = scalapack
else:
nprow = int(round(bd.comm.size**0.5))
while bd.comm.size % nprow != 0:
nprow -= 1
npcol = bd.comm.size // nprow
b = 64
p('ScaLapack grid: {0}x{1},'.format(nprow, npcol),
'block-size:', b)
bg = BlacsGrid(bd.comm, bd.comm.size, 1)
bg2 = BlacsGrid(bd.comm, nprow, npcol)
scalapack = True
else:
nprow = npcol = 1
scalapack = False
self.pt.set_positions(atoms.get_scaled_positions())
self.kpt_u[0].P_ani = None
self.allocate_arrays_for_projections(self.pt.my_atom_indices)
myslice = bd.get_slice()
pb = ProgressBar(txt)
nkpt = len(self.kpt_u)
for u, kpt in enumerate(self.kpt_u):
pb.update(u / nkpt)
npw = len(self.pd.Q_qG[kpt.q])
if scalapack:
mynpw = -(-npw // bd.comm.size)
md = BlacsDescriptor(bg, npw, npw, mynpw, npw)
md2 = BlacsDescriptor(bg2, npw, npw, b, b)
else:
md = md2 = MatrixDescriptor(npw, npw)
with self.timer('Build H and S'):
H_GG, S_GG = self.hs(ham, kpt.q, kpt.s, md)
if scalapack:
r = Redistributor(bd.comm, md, md2)
H_GG = r.redistribute(H_GG)
S_GG = r.redistribute(S_GG)
psit_nG = md2.empty(dtype=complex)
eps_n = np.empty(npw)
with self.timer('Diagonalize'):
if not scalapack:
md2.general_diagonalize_dc(H_GG, S_GG, psit_nG, eps_n,
iu=iu)
else:
md2.general_diagonalize_dc(H_GG, S_GG, psit_nG, eps_n)
del H_GG, S_GG
kpt.eps_n = eps_n[myslice].copy()
if scalapack:
md3 = BlacsDescriptor(bg, npw, npw, bd.mynbands, npw)
r = Redistributor(bd.comm, md2, md3)
psit_nG = r.redistribute(psit_nG)
kpt.psit_nG = psit_nG[:bd.mynbands].copy()
del psit_nG
with self.timer('Projections'):
self.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
kpt.f_n = None
pb.finish()
occupations.calculate(self)
def diagonalize_full_hamiltonian(self, ham, atoms, occupations, txt,
nbands=None,
scalapack=None):
if nbands is None:
nbands = self.pd.ngmin
assert nbands <= self.pd.ngmin
self.bd = bd = BandDescriptor(nbands, self.bd.comm)
if scalapack:
nprow, npcol, b = scalapack
bg = BlacsGrid(bd.comm, bd.comm.size, 1)
bg2 = BlacsGrid(bd.comm, nprow, npcol)
else:
nprow = npcol = 1
assert bd.comm.size == nprow * npcol
self.pt.set_positions(atoms.get_scaled_positions())
self.kpt_u[0].P_ani = None
self.allocate_arrays_for_projections(self.pt.my_atom_indices)
myslice = bd.get_slice()
for kpt in self.kpt_u:
npw = len(self.pd.Q_qG[kpt.q])
if scalapack:
mynpw = -(-npw // bd.comm.size)
md = BlacsDescriptor(bg, npw, npw, mynpw, npw)
md2 = BlacsDescriptor(bg2, npw, npw, b, b)
else:
md = md2 = MatrixDescriptor(npw, npw)
H_GG, S_GG = self.hs(ham, kpt.q, kpt.s, md)
if scalapack:
r = Redistributor(bd.comm, md, md2)
H_GG = r.redistribute(H_GG)
S_GG = r.redistribute(S_GG)
psit_nG = md2.empty(dtype=complex)
eps_n = np.empty(npw)
md2.general_diagonalize_dc(H_GG, S_GG, psit_nG, eps_n)
del H_GG, S_GG
kpt.eps_n = eps_n[myslice].copy()
if scalapack:
md3 = BlacsDescriptor(bg, npw, npw, bd.mynbands, npw)
r = Redistributor(bd.comm, md2, md3)
psit_nG = r.redistribute(psit_nG)
kpt.psit_nG = psit_nG[:bd.mynbands].copy()
del psit_nG
self.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
#f_n = np.zeros_like(kpt.eps_n)
#f_n[:len(kpt.f_n)] = kpt.f_n
kpt.f_n = None
occupations.calculate(self)
def diagonalize_full_hamiltonian(self, ham, atoms, occupations, log,
nbands=None, ecut=None, scalapack=None,
expert=False):
if self.dtype != complex:
raise ValueError('Your wavefunctions are not complex as '
'required by the PW diagonalization routine.\n'
'Please supply GPAW(..., dtype=complex, ...) '
'as an argument to the calculator to enforce '
'complex wavefunctions.')
if nbands is None and ecut is None:
nbands = self.pd.ngmin // self.bd.comm.size * self.bd.comm.size
elif nbands is None:
ecut /= units.Hartree
vol = abs(np.linalg.det(self.gd.cell_cv))
nbands = int(vol * ecut**1.5 * 2**0.5 / 3 / pi**2)
else:
assert nbands <= self.pd.ngmin
if expert:
iu = nbands
else:
iu = None
self.bd = bd = BandDescriptor(nbands, self.bd.comm)
log('Diagonalizing full Hamiltonian ({0} lowest bands)'.format(nbands))
log('Matrix size (min, max): {0}, {1}'.format(self.pd.ngmin,
self.pd.ngmax))
mem = 3 * self.pd.ngmax**2 * 16 / bd.comm.size / 1024**2
log('Approximate memory usage per core: {0:.3f} MB'.format(mem))
if bd.comm.size > 1:
if isinstance(scalapack, (list, tuple)):
nprow, npcol, b = scalapack
else:
nprow = int(round(bd.comm.size**0.5))
while bd.comm.size % nprow != 0:
nprow -= 1
npcol = bd.comm.size // nprow
b = 64
log('ScaLapack grid: {0}x{1},'.format(nprow, npcol),
'block-size:', b)
bg = BlacsGrid(bd.comm, bd.comm.size, 1)
bg2 = BlacsGrid(bd.comm, nprow, npcol)
scalapack = True
else:
nprow = npcol = 1
scalapack = False
self.set_positions(atoms.get_scaled_positions())
self.kpt_u[0].P_ani = None
self.allocate_arrays_for_projections(self.pt.my_atom_indices)
myslice = bd.get_slice()
pb = ProgressBar(log.fd)
nkpt = len(self.kpt_u)
for u, kpt in enumerate(self.kpt_u):
pb.update(u / nkpt)
npw = len(self.pd.Q_qG[kpt.q])
if scalapack:
mynpw = -(-npw // bd.comm.size)
md = BlacsDescriptor(bg, npw, npw, mynpw, npw)
md2 = BlacsDescriptor(bg2, npw, npw, b, b)
else:
md = md2 = MatrixDescriptor(npw, npw)
with self.timer('Build H and S'):
H_GG, S_GG = self.hs(ham, kpt.q, kpt.s, md)
if scalapack:
r = Redistributor(bd.comm, md, md2)
H_GG = r.redistribute(H_GG)
S_GG = r.redistribute(S_GG)
psit_nG = md2.empty(dtype=complex)
eps_n = np.empty(npw)
with self.timer('Diagonalize'):
if not scalapack:
md2.general_diagonalize_dc(H_GG, S_GG, psit_nG, eps_n,
iu=iu)
else:
md2.general_diagonalize_dc(H_GG, S_GG, psit_nG, eps_n)
del H_GG, S_GG
kpt.eps_n = eps_n[myslice].copy()
if scalapack:
md3 = BlacsDescriptor(bg, npw, npw, bd.mynbands, npw)
r = Redistributor(bd.comm, md2, md3)
psit_nG = r.redistribute(psit_nG)
kpt.psit_nG = psit_nG[:bd.mynbands].copy()
del psit_nG
with self.timer('Projections'):
self.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q)
kpt.f_n = None
pb.finish()
occupations.calculate(self)
return nbands