def ibz2bz(self, atoms):
"""Transform wave functions in IBZ to the full BZ."""
assert self.kd.comm.size == 1
# New k-point descriptor for full BZ:
kd = KPointDescriptor(self.kd.bzk_kc, nspins=self.nspins)
#kd.set_symmetry(atoms, self.setups, enabled=False)
kd.set_communicator(serial_comm)
self.pt = LFC(self.gd, [setup.pt_j for setup in self.setups],
kd,
dtype=self.dtype)
self.pt.set_positions(atoms.get_scaled_positions())
self.initialize_wave_functions_from_restart_file()
weight = 2.0 / kd.nspins / kd.nbzkpts
# Build new list of k-points:
kpt_u = []
for s in range(self.nspins):
for k in range(kd.nbzkpts):
# Index of symmetry related point in the IBZ
ik = self.kd.bz2ibz_k[k]
r, u = self.kd.get_rank_and_index(s, ik)
assert r == 0
kpt = self.kpt_u[u]
phase_cd = np.exp(2j * np.pi * self.gd.sdisp_cd *
kd.bzk_kc[k, :, np.newaxis])
# New k-point:
kpt2 = KPoint(weight, s, k, k, phase_cd)
kpt2.f_n = kpt.f_n / kpt.weight / kd.nbzkpts * 2 / self.nspins
kpt2.eps_n = kpt.eps_n.copy()
# Transform wave functions using symmetry operation:
Psit_nG = self.gd.collect(kpt.psit_nG)
if Psit_nG is not None:
Psit_nG = Psit_nG.copy()
for Psit_G in Psit_nG:
Psit_G[:] = self.kd.transform_wave_function(Psit_G, k)
kpt2.psit_nG = self.gd.empty(self.bd.nbands, dtype=self.dtype)
self.gd.distribute(Psit_nG, kpt2.psit_nG)
# Calculate PAW projections:
kpt2.P_ani = self.pt.dict(len(kpt.psit_nG))
self.pt.integrate(kpt2.psit_nG, kpt2.P_ani, k)
kpt_u.append(kpt2)
self.kd = kd
self.kpt_u = kpt_u
def ibz2bz(self, atoms):
"""Transform wave functions in IBZ to the full BZ."""
assert self.kd.comm.size == 1
# New k-point descriptor for full BZ:
kd = KPointDescriptor(self.kd.bzk_kc, nspins=self.nspins)
kd.set_symmetry(atoms, self.setups, usesymm=None)
kd.set_communicator(serial_comm)
self.pt = LFC(self.gd, [setup.pt_j for setup in self.setups],
kd, dtype=self.dtype)
self.pt.set_positions(atoms.get_scaled_positions())
self.initialize_wave_functions_from_restart_file()
weight = 2.0 / kd.nspins / kd.nbzkpts
# Build new list of k-points:
kpt_u = []
for s in range(self.nspins):
for k in range(kd.nbzkpts):
# Index of symmetry related point in the IBZ
ik = self.kd.bz2ibz_k[k]
r, u = self.kd.get_rank_and_index(s, ik)
assert r == 0
kpt = self.kpt_u[u]
phase_cd = np.exp(2j * np.pi * self.gd.sdisp_cd *
kd.bzk_kc[k, :, np.newaxis])
# New k-point:
kpt2 = KPoint(weight, s, k, k, phase_cd)
kpt2.f_n = kpt.f_n / kpt.weight / kd.nbzkpts * 2 / self.nspins
kpt2.eps_n = kpt.eps_n.copy()
# Transform wave functions using symmetry operation:
Psit_nG = self.gd.collect(kpt.psit_nG)
if Psit_nG is not None:
Psit_nG = Psit_nG.copy()
for Psit_G in Psit_nG:
Psit_G[:] = self.kd.transform_wave_function(Psit_G, k)
kpt2.psit_nG = self.gd.empty(self.bd.nbands, dtype=self.dtype)
self.gd.distribute(Psit_nG, kpt2.psit_nG)
# Calculate PAW projections:
kpt2.P_ani = self.pt.dict(len(kpt.psit_nG))
self.pt.integrate(kpt2.psit_nG, kpt2.P_ani, k)
kpt_u.append(kpt2)
self.kd = kd
self.kpt_u = kpt_u
def initialize(self, atoms=None):
"""Inexpensive initialization."""
if atoms is None:
atoms = self.atoms
else:
# Save the state of the atoms:
self.atoms = atoms.copy()
par = self.input_parameters
world = par.communicator
if world is None:
world = mpi.world
elif hasattr(world, 'new_communicator'):
# Check for whether object has correct type already
#
# Using isinstance() is complicated because of all the
# combinations, serial/parallel/debug...
pass
else:
# world should be a list of ranks:
world = mpi.world.new_communicator(np.asarray(world))
self.wfs.world = world
self.set_text(par.txt, par.verbose)
natoms = len(atoms)
cell_cv = atoms.get_cell() / Bohr
pbc_c = atoms.get_pbc()
Z_a = atoms.get_atomic_numbers()
magmom_av = atoms.get_initial_magnetic_moments()
# Generate new xc functional only when it is reset by set
if self.hamiltonian is None or self.hamiltonian.xc is None:
if isinstance(par.xc, str):
xc = XC(par.xc)
else:
xc = par.xc
else:
xc = self.hamiltonian.xc
mode = par.mode
if xc.orbital_dependent and mode == 'lcao':
raise NotImplementedError('LCAO mode does not support '
'orbital-dependent XC functionals.')
if mode == 'pw':
mode = PW()
if mode == 'fd' and par.usefractrans:
raise NotImplementedError('FD mode does not support '
'fractional translations.')
if mode == 'lcao' and par.usefractrans:
raise Warning('Fractional translations have not been tested '
'with LCAO mode. Use with care!')
if par.realspace is None:
realspace = not isinstance(mode, PW)
else:
realspace = par.realspace
if isinstance(mode, PW):
assert not realspace
if par.gpts is not None:
N_c = np.array(par.gpts)
else:
h = par.h
if h is not None:
h /= Bohr
N_c = get_number_of_grid_points(cell_cv, h, mode, realspace)
if par.filter is None and not isinstance(mode, PW):
gamma = 1.6
hmax = ((np.linalg.inv(cell_cv)**2).sum(0)**-0.5 / N_c).max()
def filter(rgd, rcut, f_r, l=0):
gcut = np.pi / hmax - 2 / rcut / gamma
f_r[:] = rgd.filter(f_r, rcut * gamma, gcut, l)
else:
filter = par.filter
setups = Setups(Z_a, par.setups, par.basis, par.lmax, xc,
filter, world)
if magmom_av.ndim == 1:
collinear = True
magmom_av, magmom_a = np.zeros((natoms, 3)), magmom_av
magmom_av[:, 2] = magmom_a
else:
collinear = False
magnetic = magmom_av.any()
spinpol = par.spinpol
if par.hund:
if natoms != 1:
raise ValueError('hund=True arg only valid for single atoms!')
spinpol = True
magmom_av[0] = (0, 0, setups[0].get_hunds_rule_moment(par.charge))
if spinpol is None:
spinpol = magnetic
elif magnetic and not spinpol:
raise ValueError('Non-zero initial magnetic moment for a ' +
'spin-paired calculation!')
if collinear:
nspins = 1 + int(spinpol)
ncomp = 1
else:
nspins = 1
ncomp = 2
# K-point descriptor
bzkpts_kc = kpts2ndarray(par.kpts, self.atoms)
kd = KPointDescriptor(bzkpts_kc, nspins, collinear, par.usefractrans)
width = par.width
if width is None:
if pbc_c.any():
width = 0.1 # eV
else:
width = 0.0
else:
assert par.occupations is None
if hasattr(self, 'time') or par.dtype == complex:
dtype = complex
else:
if kd.gamma:
dtype = float
else:
dtype = complex
## rbw: If usefractrans=True, kd.set_symmetry might overwrite N_c.
## This is necessary, because N_c must be dividable by 1/(fractional translation),
## f.e. fractional translations of a grid point must land on a grid point.
N_c = kd.set_symmetry(atoms, setups, magmom_av, par.usesymm, N_c, world)
nao = setups.nao
nvalence = setups.nvalence - par.charge
M_v = magmom_av.sum(0)
M = np.dot(M_v, M_v)**0.5
nbands = par.nbands
if nbands is None:
nbands = 0
for setup in setups:
nbands_from_atom = setup.get_default_nbands()
# Any obscure setup errors?
if nbands_from_atom < -(-setup.Nv // 2):
raise ValueError('Bad setup: This setup requests %d'
' bands but has %d electrons.'
% (nbands_from_atom, setup.Nv))
nbands += nbands_from_atom
nbands = min(nao, nbands)
elif nbands > nao and mode == 'lcao':
raise ValueError('Too many bands for LCAO calculation: '
'%d bands and only %d atomic orbitals!' %
(nbands, nao))
if nvalence < 0:
raise ValueError(
'Charge %f is not possible - not enough valence electrons' %
par.charge)
if nbands <= 0:
nbands = int(nvalence + M + 0.5) // 2 + (-nbands)
if nvalence > 2 * nbands:
raise ValueError('Too few bands! Electrons: %f, bands: %d'
% (nvalence, nbands))
nbands *= ncomp
if par.width is not None:
self.text('**NOTE**: please start using '
'occupations=FermiDirac(width).')
if par.fixmom:
self.text('**NOTE**: please start using '
'occupations=FermiDirac(width, fixmagmom=True).')
if self.occupations is None:
if par.occupations is None:
# Create object for occupation numbers:
self.occupations = occupations.FermiDirac(width, par.fixmom)
else:
self.occupations = par.occupations
self.occupations.magmom = M_v[2]
cc = par.convergence
if mode == 'lcao':
niter_fixdensity = 0
else:
niter_fixdensity = None
if self.scf is None:
self.scf = SCFLoop(
cc['eigenstates'] / Hartree**2 * nvalence,
cc['energy'] / Hartree * max(nvalence, 1),
cc['density'] * nvalence,
par.maxiter, par.fixdensity,
niter_fixdensity)
parsize_kpt = par.parallel['kpt']
parsize_domain = par.parallel['domain']
parsize_bands = par.parallel['band']
if not realspace:
pbc_c = np.ones(3, bool)
if not self.wfs:
if parsize_domain == 'domain only': # XXX this was silly!
parsize_domain = world.size
parallelization = mpi.Parallelization(world,
nspins * kd.nibzkpts)
ndomains = None
if parsize_domain is not None:
ndomains = np.prod(parsize_domain)
if isinstance(mode, PW):
if ndomains > 1:
raise ValueError('Planewave mode does not support '
'domain decomposition.')
ndomains = 1
parallelization.set(kpt=parsize_kpt,
domain=ndomains,
band=parsize_bands)
domain_comm, kpt_comm, band_comm = \
parallelization.build_communicators()
#domain_comm, kpt_comm, band_comm = mpi.distribute_cpus(
# parsize_domain, parsize_bands,
# nspins, kd.nibzkpts, world, par.idiotproof, mode)
kd.set_communicator(kpt_comm)
parstride_bands = par.parallel['stridebands']
# Unfortunately we need to remember that we adjusted the
# number of bands so we can print a warning if it differs
# from the number specified by the user. (The number can
# be inferred from the input parameters, but it's tricky
# because we allow negative numbers)
self.nbands_parallelization_adjustment = -nbands % band_comm.size
nbands += self.nbands_parallelization_adjustment
# I would like to give the following error message, but apparently
# there are cases, e.g. gpaw/test/gw_ppa.py, which involve
# nbands > nao and are supposed to work that way.
#if nbands > nao:
# raise ValueError('Number of bands %d adjusted for band '
# 'parallelization %d exceeds number of atomic '
# 'orbitals %d. This problem can be fixed '
# 'by reducing the number of bands a bit.'
# % (nbands, band_comm.size, nao))
bd = BandDescriptor(nbands, band_comm, parstride_bands)
if (self.density is not None and
self.density.gd.comm.size != domain_comm.size):
# Domain decomposition has changed, so we need to
# reinitialize density and hamiltonian:
if par.fixdensity:
raise RuntimeError('Density reinitialization conflict ' +
'with "fixdensity" - specify domain decomposition.')
self.density = None
self.hamiltonian = None
# Construct grid descriptor for coarse grids for wave functions:
gd = self.grid_descriptor_class(N_c, cell_cv, pbc_c,
domain_comm, parsize_domain)
# do k-point analysis here? XXX
args = (gd, nvalence, setups, bd, dtype, world, kd, self.timer)
if par.parallel['sl_auto']:
# Choose scalapack parallelization automatically
for key, val in par.parallel.items():
if (key.startswith('sl_') and key != 'sl_auto'
and val is not None):
raise ValueError("Cannot use 'sl_auto' together "
"with '%s'" % key)
max_scalapack_cpus = bd.comm.size * gd.comm.size
nprow = max_scalapack_cpus
npcol = 1
# Get a sort of reasonable number of columns/rows
while npcol < nprow and nprow % 2 == 0:
npcol *= 2
nprow //= 2
assert npcol * nprow == max_scalapack_cpus
# ScaLAPACK creates trouble if there aren't at least a few
# whole blocks; choose block size so there will always be
# several blocks. This will crash for small test systems,
# but so will ScaLAPACK in any case
blocksize = min(-(-nbands // 4), 64)
sl_default = (nprow, npcol, blocksize)
else:
sl_default = par.parallel['sl_default']
if mode == 'lcao':
# Layouts used for general diagonalizer
sl_lcao = par.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = sl_default
lcaoksl = get_KohnSham_layouts(sl_lcao, 'lcao',
gd, bd, dtype,
nao=nao, timer=self.timer)
if collinear:
self.wfs = LCAOWaveFunctions(lcaoksl, *args)
else:
from gpaw.xc.noncollinear import \
NonCollinearLCAOWaveFunctions
self.wfs = NonCollinearLCAOWaveFunctions(lcaoksl, *args)
elif mode == 'fd' or isinstance(mode, PW):
# buffer_size keyword only relevant for fdpw
buffer_size = par.parallel['buffer_size']
# Layouts used for diagonalizer
sl_diagonalize = par.parallel['sl_diagonalize']
if sl_diagonalize is None:
sl_diagonalize = sl_default
diagksl = get_KohnSham_layouts(sl_diagonalize, 'fd',
gd, bd, dtype,
buffer_size=buffer_size,
timer=self.timer)
# Layouts used for orthonormalizer
sl_inverse_cholesky = par.parallel['sl_inverse_cholesky']
if sl_inverse_cholesky is None:
sl_inverse_cholesky = sl_default
if sl_inverse_cholesky != sl_diagonalize:
message = 'sl_inverse_cholesky != sl_diagonalize ' \
'is not implemented.'
raise NotImplementedError(message)
orthoksl = get_KohnSham_layouts(sl_inverse_cholesky, 'fd',
gd, bd, dtype,
buffer_size=buffer_size,
timer=self.timer)
# Use (at most) all available LCAO for initialization
lcaonbands = min(nbands, nao)
try:
lcaobd = BandDescriptor(lcaonbands, band_comm,
parstride_bands)
except RuntimeError:
initksl = None
else:
# Layouts used for general diagonalizer
# (LCAO initialization)
sl_lcao = par.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = sl_default
initksl = get_KohnSham_layouts(sl_lcao, 'lcao',
gd, lcaobd, dtype,
nao=nao,
timer=self.timer)
if hasattr(self, 'time'):
assert mode == 'fd'
from gpaw.tddft import TimeDependentWaveFunctions
self.wfs = TimeDependentWaveFunctions(par.stencils[0],
diagksl, orthoksl, initksl, gd, nvalence, setups,
bd, world, kd, self.timer)
elif mode == 'fd':
self.wfs = FDWaveFunctions(par.stencils[0], diagksl,
orthoksl, initksl, *args)
else:
# Planewave basis:
self.wfs = mode(diagksl, orthoksl, initksl, *args)
else:
self.wfs = mode(self, *args)
else:
self.wfs.set_setups(setups)
if not self.wfs.eigensolver:
# Number of bands to converge:
nbands_converge = cc['bands']
if nbands_converge == 'all':
nbands_converge = nbands
elif nbands_converge != 'occupied':
assert isinstance(nbands_converge, int)
if nbands_converge < 0:
nbands_converge += nbands
eigensolver = get_eigensolver(par.eigensolver, mode,
par.convergence)
eigensolver.nbands_converge = nbands_converge
# XXX Eigensolver class doesn't define an nbands_converge property
if isinstance(xc, SIC):
eigensolver.blocksize = 1
self.wfs.set_eigensolver(eigensolver)
if self.density is None:
gd = self.wfs.gd
if par.stencils[1] != 9:
# Construct grid descriptor for fine grids for densities
# and potentials:
finegd = gd.refine()
else:
# Special case (use only coarse grid):
finegd = gd
if realspace:
self.density = RealSpaceDensity(
gd, finegd, nspins, par.charge + setups.core_charge,
collinear, par.stencils[1])
else:
self.density = ReciprocalSpaceDensity(
gd, finegd, nspins, par.charge + setups.core_charge,
collinear)
self.density.initialize(setups, self.timer, magmom_av, par.hund)
self.density.set_mixer(par.mixer)
if self.hamiltonian is None:
gd, finegd = self.density.gd, self.density.finegd
if realspace:
self.hamiltonian = RealSpaceHamiltonian(
gd, finegd, nspins, setups, self.timer, xc, par.external,
collinear, par.poissonsolver, par.stencils[1], world)
else:
self.hamiltonian = ReciprocalSpaceHamiltonian(
gd, finegd,
self.density.pd2, self.density.pd3,
nspins, setups, self.timer, xc, par.external,
collinear, world)
xc.initialize(self.density, self.hamiltonian, self.wfs,
self.occupations)
self.text()
self.print_memory_estimate(self.txt, maxdepth=memory_estimate_depth)
self.txt.flush()
self.timer.print_info(self)
if dry_run:
self.dry_run()
self.initialized = True
class UTGaussianWavefunctionSetup(UTDomainParallelSetup):
__doc__ = UTDomainParallelSetup.__doc__ + """
The pseudo wavefunctions are moving gaussians centered around each atom."""
allocated = False
dtype = None
# Default arguments for scaled Gaussian wave
_sigma0 = 2.0 #0.75
_k0_c = 2*np.pi*np.array([1/5., 1/3., 0.])
def setUp(self):
UTDomainParallelSetup.setUp(self)
for virtvar in ['dtype']:
assert getattr(self,virtvar) is not None, 'Virtual "%s"!' % virtvar
# Create randomized atoms
self.atoms = create_random_atoms(self.gd, 5) # also tested: 10xNH3/BDA
# XXX DEBUG START
if False:
from ase import view
view(self.atoms*(1+2*self.gd.pbc_c))
# XXX DEBUG END
# Do we agree on the atomic positions?
pos_ac = self.atoms.get_positions()
pos_rac = np.empty((world.size,)+pos_ac.shape, pos_ac.dtype)
world.all_gather(pos_ac, pos_rac)
if (pos_rac-pos_rac[world.rank,...][np.newaxis,...]).any():
raise RuntimeError('Discrepancy in atomic positions detected.')
# Create setups for atoms
self.Z_a = self.atoms.get_atomic_numbers()
self.setups = Setups(self.Z_a, p.setups, p.basis,
p.lmax, xc)
# K-point descriptor
bzk_kc = np.array([[0, 0, 0]], dtype=float)
self.kd = KPointDescriptor(bzk_kc, 1)
self.kd.set_symmetry(self.atoms, self.setups)
self.kd.set_communicator(self.kpt_comm)
# Create gamma-point dummy wavefunctions
self.wfs = FDWFS(self.gd, self.bd, self.kd, self.setups,
self.block_comm, self.dtype)
spos_ac = self.atoms.get_scaled_positions() % 1.0
self.wfs.set_positions(spos_ac)
self.pt = self.wfs.pt # XXX shortcut
## Also create pseudo partial waveves
#from gpaw.lfc import LFC
#self.phit = LFC(self.gd, [setup.phit_j for setup in self.setups], \
# self.kpt_comm, dtype=self.dtype)
#self.phit.set_positions(spos_ac)
self.r_cG = None
self.buf_G = None
self.psit_nG = None
self.allocate()
def tearDown(self):
UTDomainParallelSetup.tearDown(self)
del self.r_cG, self.buf_G, self.psit_nG
del self.pt, self.setups, self.atoms
self.allocated = False
def allocate(self):
self.r_cG = self.gd.get_grid_point_coordinates()
cell_cv = self.atoms.get_cell() / Bohr
assert np.abs(cell_cv-self.gd.cell_cv).max() < 1e-9
center_c = 0.5*cell_cv.diagonal()
self.buf_G = self.gd.empty(dtype=self.dtype)
self.psit_nG = self.gd.empty(self.bd.mynbands, dtype=self.dtype)
for myn,psit_G in enumerate(self.psit_nG):
n = self.bd.global_index(myn)
psit_G[:] = self.get_scaled_gaussian_wave(center_c, scale=10+2j*n)
k_c = 2*np.pi*np.array([1/2., -1/7., 0.])
for pos_c in self.atoms.get_positions() / Bohr:
sigma = self._sigma0/(1+np.sum(pos_c**2))**0.5
psit_G += self.get_scaled_gaussian_wave(pos_c, sigma, k_c, n+5j)
self.allocated = True
def get_scaled_gaussian_wave(self, pos_c, sigma=None, k_c=None, scale=None):
if sigma is None:
sigma = self._sigma0
if k_c is None:
k_c = self._k0_c
if scale is None:
A = None
else:
# 4*pi*int(exp(-r^2/(2*w^2))^2*r^2, r=0...infinity)= w^3*pi^(3/2)
# = scale/A^2 -> A = scale*(sqrt(Pi)*w)^(-3/2) hence int -> scale^2
A = scale/(sigma*(np.pi)**0.5)**1.5
return gaussian_wave(self.r_cG, pos_c, sigma, k_c, A, self.dtype, self.buf_G)
def check_and_plot(self, P_ani, P0_ani, digits, keywords=''):
# Collapse into viewable matrices
P_In = self.wfs.collect_projections(P_ani)
P0_In = self.wfs.collect_projections(P0_ani)
# Construct fingerprint of input matrices for comparison
fingerprint = np.array([md5_array(P_In, numeric=True),
md5_array(P0_In, numeric=True)])
# Compare fingerprints across all processors
fingerprints = np.empty((world.size, 2), np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
# If assertion fails, catch temporarily while plotting, then re-raise
try:
self.assertAlmostEqual(np.abs(P_In-P0_In).max(), 0, digits)
except AssertionError:
if world.rank == 0 and mpl is not None:
from matplotlib.figure import Figure
fig = Figure()
ax = fig.add_axes([0.0, 0.1, 1.0, 0.83])
ax.set_title(self.__class__.__name__)
im = ax.imshow(np.abs(P_In-P0_In), interpolation='nearest')
fig.colorbar(im)
fig.text(0.5, 0.05, 'Keywords: ' + keywords, \
horizontalalignment='center', verticalalignment='top')
from matplotlib.backends.backend_agg import FigureCanvasAgg
img = 'ut_invops_%s_%s.png' % (self.__class__.__name__, \
'_'.join(keywords.split(',')))
FigureCanvasAgg(fig).print_figure(img.lower(), dpi=90)
raise
# =================================
def test_projection_linearity(self):
kpt = self.wfs.kpt_u[0]
Q_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(self.psit_nG, Q_ani, q=kpt.q)
for Q_ni in Q_ani.values():
self.assertTrue(Q_ni.dtype == self.dtype)
P0_ani = dict([(a,Q_ni.copy()) for a,Q_ni in Q_ani.items()])
self.pt.add(self.psit_nG, Q_ani, q=kpt.q)
self.pt.integrate(self.psit_nG, P0_ani, q=kpt.q)
#rank_a = self.gd.get_ranks_from_positions(spos_ac)
#my_atom_indices = np.argwhere(self.gd.comm.rank == rank_a).ravel()
# ~ a ~ a'
#TODO XXX should fix PairOverlap-ish stuff for < p | phi > overlaps
# i i'
#spos_ac = self.pt.spos_ac # NewLFC doesn't have this
spos_ac = self.atoms.get_scaled_positions() % 1.0
gpo = GridPairOverlap(self.gd, self.setups)
B_aa = gpo.calculate_overlaps(spos_ac, self.pt)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
P_ani = dict([(a,Q_ni.copy()) for a,Q_ni in Q_ani.items()])
for a1 in range(len(self.atoms)):
if a1 in P_ani.keys():
P_ni = P_ani[a1]
else:
# Atom a1 is not in domain so allocate a temporary buffer
P_ni = np.zeros((self.bd.mynbands,self.setups[a1].ni,),
dtype=self.dtype)
for a2, Q_ni in Q_ani.items():
# B_aa are the projector overlaps across atomic pairs
B_ii = gpo.extract_atomic_pair_matrix(B_aa, a1, a2)
P_ni += np.dot(Q_ni, B_ii.T) #sum over a2 and last i in B_ii
self.gd.comm.sum(P_ni)
self.check_and_plot(P_ani, P0_ani, 8, 'projection,linearity')
def test_extrapolate_overlap(self):
kpt = self.wfs.kpt_u[0]
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
P_ani = ppo.apply(self.psit_nG, work_nG, self.wfs, kpt, \
calculate_P_ani=True, extrapolate_P_ani=True)
P0_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(work_nG, P0_ani, kpt.q)
del work_nG
self.check_and_plot(P_ani, P0_ani, 11, 'extrapolate,overlap')
def test_extrapolate_inverse(self):
kpt = self.wfs.kpt_u[0]
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
P_ani = ppo.apply_inverse(self.psit_nG, work_nG, self.wfs, kpt, \
calculate_P_ani=True, extrapolate_P_ani=True)
P0_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(work_nG, P0_ani, kpt.q)
del work_nG
self.check_and_plot(P_ani, P0_ani, 11, 'extrapolate,inverse')
def test_overlap_inverse_after(self):
kpt = self.wfs.kpt_u[0]
kpt.P_ani = self.pt.dict(self.bd.mynbands)
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
self.pt.integrate(self.psit_nG, kpt.P_ani, kpt.q)
P0_ani = dict([(a,P_ni.copy()) for a,P_ni in kpt.P_ani.items()])
ppo.apply(self.psit_nG, work_nG, self.wfs, kpt, calculate_P_ani=False)
res_nG = np.empty_like(self.psit_nG)
ppo.apply_inverse(work_nG, res_nG, self.wfs, kpt, calculate_P_ani=True)
del work_nG
P_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(res_nG, P_ani, kpt.q)
abserr = np.empty(1, dtype=float)
for n in range(self.nbands):
band_rank, myn = self.bd.who_has(n)
if band_rank == self.bd.comm.rank:
abserr[:] = np.abs(self.psit_nG[myn] - res_nG[myn]).max()
self.gd.comm.max(abserr)
self.bd.comm.broadcast(abserr, band_rank)
self.assertAlmostEqual(abserr.item(), 0, 10)
self.check_and_plot(P_ani, P0_ani, 10, 'overlap,inverse,after')
def test_overlap_inverse_before(self):
kpt = self.wfs.kpt_u[0]
kpt.P_ani = self.pt.dict(self.bd.mynbands)
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
self.pt.integrate(self.psit_nG, kpt.P_ani, kpt.q)
P0_ani = dict([(a,P_ni.copy()) for a,P_ni in kpt.P_ani.items()])
ppo.apply_inverse(self.psit_nG, work_nG, self.wfs, kpt, calculate_P_ani=False)
res_nG = np.empty_like(self.psit_nG)
ppo.apply(work_nG, res_nG, self.wfs, kpt, calculate_P_ani=True)
del work_nG
P_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(res_nG, P_ani, kpt.q)
abserr = np.empty(1, dtype=float)
for n in range(self.nbands):
band_rank, myn = self.bd.who_has(n)
if band_rank == self.bd.comm.rank:
abserr[:] = np.abs(self.psit_nG[myn] - res_nG[myn]).max()
self.gd.comm.max(abserr)
self.bd.comm.broadcast(abserr, band_rank)
self.assertAlmostEqual(abserr.item(), 0, 10)
self.check_and_plot(P_ani, P0_ani, 10, 'overlap,inverse,before')
class UTGaussianWavefunctionSetup(UTDomainParallelSetup):
__doc__ = UTDomainParallelSetup.__doc__ + """
The pseudo wavefunctions are moving gaussians centered around each atom."""
allocated = False
dtype = None
# Default arguments for scaled Gaussian wave
_sigma0 = 2.0 #0.75
_k0_c = 2*np.pi*np.array([1/5., 1/3., 0.])
def setUp(self):
UTDomainParallelSetup.setUp(self)
for virtvar in ['dtype']:
assert getattr(self,virtvar) is not None, 'Virtual "%s"!' % virtvar
# Create randomized atoms
self.atoms = create_random_atoms(self.gd, 5) # also tested: 10xNH3/BDA
# XXX DEBUG START
if False:
from ase import view
view(self.atoms*(1+2*self.gd.pbc_c))
# XXX DEBUG END
# Do we agree on the atomic positions?
pos_ac = self.atoms.get_positions()
pos_rac = np.empty((world.size,)+pos_ac.shape, pos_ac.dtype)
world.all_gather(pos_ac, pos_rac)
if (pos_rac-pos_rac[world.rank,...][np.newaxis,...]).any():
raise RuntimeError('Discrepancy in atomic positions detected.')
# Create setups for atoms
self.Z_a = self.atoms.get_atomic_numbers()
self.setups = Setups(self.Z_a, p.setups, p.basis,
p.lmax, xc)
# K-point descriptor
bzk_kc = np.array([[0, 0, 0]], dtype=float)
self.kd = KPointDescriptor(bzk_kc, 1)
self.kd.set_symmetry(self.atoms, self.setups, usesymm=True)
self.kd.set_communicator(self.kpt_comm)
# Create gamma-point dummy wavefunctions
self.wfs = FDWFS(self.gd, self.bd, self.kd, self.setups,
self.dtype)
spos_ac = self.atoms.get_scaled_positions() % 1.0
self.wfs.set_positions(spos_ac)
self.pt = self.wfs.pt # XXX shortcut
## Also create pseudo partial waveves
#from gpaw.lfc import LFC
#self.phit = LFC(self.gd, [setup.phit_j for setup in self.setups], \
# self.kpt_comm, dtype=self.dtype)
#self.phit.set_positions(spos_ac)
self.r_cG = None
self.buf_G = None
self.psit_nG = None
self.allocate()
def tearDown(self):
UTDomainParallelSetup.tearDown(self)
del self.r_cG, self.buf_G, self.psit_nG
del self.pt, self.setups, self.atoms
self.allocated = False
def allocate(self):
self.r_cG = self.gd.get_grid_point_coordinates()
cell_cv = self.atoms.get_cell() / Bohr
assert np.abs(cell_cv-self.gd.cell_cv).max() < 1e-9
center_c = 0.5*cell_cv.diagonal()
self.buf_G = self.gd.empty(dtype=self.dtype)
self.psit_nG = self.gd.empty(self.bd.mynbands, dtype=self.dtype)
for myn,psit_G in enumerate(self.psit_nG):
n = self.bd.global_index(myn)
psit_G[:] = self.get_scaled_gaussian_wave(center_c, scale=10+2j*n)
k_c = 2*np.pi*np.array([1/2., -1/7., 0.])
for pos_c in self.atoms.get_positions() / Bohr:
sigma = self._sigma0/(1+np.sum(pos_c**2))**0.5
psit_G += self.get_scaled_gaussian_wave(pos_c, sigma, k_c, n+5j)
self.allocated = True
def get_scaled_gaussian_wave(self, pos_c, sigma=None, k_c=None, scale=None):
if sigma is None:
sigma = self._sigma0
if k_c is None:
k_c = self._k0_c
if scale is None:
A = None
else:
# 4*pi*int(exp(-r^2/(2*w^2))^2*r^2, r=0...infinity)= w^3*pi^(3/2)
# = scale/A^2 -> A = scale*(sqrt(Pi)*w)^(-3/2) hence int -> scale^2
A = scale/(sigma*(np.pi)**0.5)**1.5
return gaussian_wave(self.r_cG, pos_c, sigma, k_c, A, self.dtype, self.buf_G)
def check_and_plot(self, P_ani, P0_ani, digits, keywords=''):
# Collapse into viewable matrices
P_In = self.wfs.collect_projections(P_ani)
P0_In = self.wfs.collect_projections(P0_ani)
# Construct fingerprint of input matrices for comparison
fingerprint = np.array([md5_array(P_In, numeric=True),
md5_array(P0_In, numeric=True)])
# Compare fingerprints across all processors
fingerprints = np.empty((world.size, 2), np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
# If assertion fails, catch temporarily while plotting, then re-raise
try:
self.assertAlmostEqual(np.abs(P_In-P0_In).max(), 0, digits)
except AssertionError:
if world.rank == 0 and mpl is not None:
from matplotlib.figure import Figure
fig = Figure()
ax = fig.add_axes([0.0, 0.1, 1.0, 0.83])
ax.set_title(self.__class__.__name__)
im = ax.imshow(np.abs(P_In-P0_In), interpolation='nearest')
fig.colorbar(im)
fig.text(0.5, 0.05, 'Keywords: ' + keywords, \
horizontalalignment='center', verticalalignment='top')
from matplotlib.backends.backend_agg import FigureCanvasAgg
img = 'ut_invops_%s_%s.png' % (self.__class__.__name__, \
'_'.join(keywords.split(',')))
FigureCanvasAgg(fig).print_figure(img.lower(), dpi=90)
raise
# =================================
def test_projection_linearity(self):
kpt = self.wfs.kpt_u[0]
Q_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(self.psit_nG, Q_ani, q=kpt.q)
for Q_ni in Q_ani.values():
self.assertTrue(Q_ni.dtype == self.dtype)
P0_ani = dict([(a,Q_ni.copy()) for a,Q_ni in Q_ani.items()])
self.pt.add(self.psit_nG, Q_ani, q=kpt.q)
self.pt.integrate(self.psit_nG, P0_ani, q=kpt.q)
#rank_a = self.gd.get_ranks_from_positions(spos_ac)
#my_atom_indices = np.argwhere(self.gd.comm.rank == rank_a).ravel()
# ~ a ~ a'
#TODO XXX should fix PairOverlap-ish stuff for < p | phi > overlaps
# i i'
#spos_ac = self.pt.spos_ac # NewLFC doesn't have this
spos_ac = self.atoms.get_scaled_positions() % 1.0
gpo = GridPairOverlap(self.gd, self.setups)
B_aa = gpo.calculate_overlaps(spos_ac, self.pt)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
P_ani = dict([(a,Q_ni.copy()) for a,Q_ni in Q_ani.items()])
for a1 in range(len(self.atoms)):
if a1 in P_ani.keys():
P_ni = P_ani[a1]
else:
# Atom a1 is not in domain so allocate a temporary buffer
P_ni = np.zeros((self.bd.mynbands,self.setups[a1].ni,),
dtype=self.dtype)
for a2, Q_ni in Q_ani.items():
# B_aa are the projector overlaps across atomic pairs
B_ii = gpo.extract_atomic_pair_matrix(B_aa, a1, a2)
P_ni += np.dot(Q_ni, B_ii.T) #sum over a2 and last i in B_ii
self.gd.comm.sum(P_ni)
self.check_and_plot(P_ani, P0_ani, 8, 'projection,linearity')
def test_extrapolate_overlap(self):
kpt = self.wfs.kpt_u[0]
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
P_ani = ppo.apply(self.psit_nG, work_nG, self.wfs, kpt, \
calculate_P_ani=True, extrapolate_P_ani=True)
P0_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(work_nG, P0_ani, kpt.q)
del work_nG
self.check_and_plot(P_ani, P0_ani, 11, 'extrapolate,overlap')
def test_extrapolate_inverse(self):
kpt = self.wfs.kpt_u[0]
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
P_ani = ppo.apply_inverse(self.psit_nG, work_nG, self.wfs, kpt, \
calculate_P_ani=True, extrapolate_P_ani=True)
P0_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(work_nG, P0_ani, kpt.q)
del work_nG
self.check_and_plot(P_ani, P0_ani, 11, 'extrapolate,inverse')
def test_overlap_inverse_after(self):
kpt = self.wfs.kpt_u[0]
kpt.P_ani = self.pt.dict(self.bd.mynbands)
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
self.pt.integrate(self.psit_nG, kpt.P_ani, kpt.q)
P0_ani = dict([(a,P_ni.copy()) for a,P_ni in kpt.P_ani.items()])
ppo.apply(self.psit_nG, work_nG, self.wfs, kpt, calculate_P_ani=False)
res_nG = np.empty_like(self.psit_nG)
ppo.apply_inverse(work_nG, res_nG, self.wfs, kpt, calculate_P_ani=True)
del work_nG
P_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(res_nG, P_ani, kpt.q)
abserr = np.empty(1, dtype=float)
for n in range(self.nbands):
band_rank, myn = self.bd.who_has(n)
if band_rank == self.bd.comm.rank:
abserr[:] = np.abs(self.psit_nG[myn] - res_nG[myn]).max()
self.gd.comm.max(abserr)
self.bd.comm.broadcast(abserr, band_rank)
self.assertAlmostEqual(abserr.item(), 0, 10)
self.check_and_plot(P_ani, P0_ani, 10, 'overlap,inverse,after')
def test_overlap_inverse_before(self):
kpt = self.wfs.kpt_u[0]
kpt.P_ani = self.pt.dict(self.bd.mynbands)
ppo = ProjectorPairOverlap(self.wfs, self.atoms)
# Compare fingerprints across all processors
fingerprint = np.array([md5_array(ppo.B_aa, numeric=True)])
fingerprints = np.empty(world.size, np.int64)
world.all_gather(fingerprint, fingerprints)
if fingerprints.ptp(0).any():
raise RuntimeError('Distributed matrices are not identical!')
work_nG = np.empty_like(self.psit_nG)
self.pt.integrate(self.psit_nG, kpt.P_ani, kpt.q)
P0_ani = dict([(a,P_ni.copy()) for a,P_ni in kpt.P_ani.items()])
ppo.apply_inverse(self.psit_nG, work_nG, self.wfs, kpt, calculate_P_ani=False)
res_nG = np.empty_like(self.psit_nG)
ppo.apply(work_nG, res_nG, self.wfs, kpt, calculate_P_ani=True)
del work_nG
P_ani = self.pt.dict(self.bd.mynbands)
self.pt.integrate(res_nG, P_ani, kpt.q)
abserr = np.empty(1, dtype=float)
for n in range(self.nbands):
band_rank, myn = self.bd.who_has(n)
if band_rank == self.bd.comm.rank:
abserr[:] = np.abs(self.psit_nG[myn] - res_nG[myn]).max()
self.gd.comm.max(abserr)
self.bd.comm.broadcast(abserr, band_rank)
self.assertAlmostEqual(abserr.item(), 0, 10)
self.check_and_plot(P_ani, P0_ani, 10, 'overlap,inverse,before')
class PhononCalculator:
"""This class defines the interface for phonon calculations."""
def __init__(self,
calc,
gamma=True,
symmetry=False,
e_ph=False,
communicator=serial_comm):
"""Inititialize class with a list of atoms.
The atoms object must contain a converged ground-state calculation.
The set of q-vectors in which the dynamical matrix will be calculated
is determined from the ``symmetry`` kwarg. For now, only time-reversal
symmetry is used to generate the irrecducible BZ.
Add a little note on parallelization strategy here.
Parameters
----------
calc: str or Calculator
Calculator containing a ground-state calculation.
gamma: bool
Gamma-point calculation with respect to the q-vector of the
dynamical matrix. When ``False``, the Monkhorst-Pack grid from the
ground-state calculation is used.
symmetry: bool
Use symmetries to reduce the q-vectors of the dynamcial matrix
(None, False or True). The different options are equivalent to the
old style options in a ground-state calculation (see usesymm).
e_ph: bool
Save the derivative of the effective potential.
communicator: Communicator
Communicator for parallelization over k-points and real-space
domain.
"""
# XXX
assert symmetry in [None, False], "Spatial symmetries not allowed yet"
if isinstance(calc, str):
self.calc = GPAW(calc, communicator=serial_comm, txt=None)
else:
self.calc = calc
cell_cv = self.calc.atoms.get_cell()
setups = self.calc.wfs.setups
# XXX - no clue how to get magmom - ignore it for the moment
# m_av = magmom_av.round(decimals=3) # round off
# id_a = zip(setups.id_a, *m_av.T)
id_a = setups.id_a
if symmetry is None:
self.symmetry = Symmetry(id_a,
cell_cv,
point_group=False,
time_reversal=False)
else:
self.symmetry = Symmetry(id_a,
cell_cv,
point_group=False,
time_reversal=True)
# Make sure localized functions are initialized
self.calc.set_positions()
# Note that this under some circumstances (e.g. when called twice)
# allocates a new array for the P_ani coefficients !!
# Store useful objects
self.atoms = self.calc.get_atoms()
# Get rid of ``calc`` attribute
self.atoms.calc = None
# Boundary conditions
pbc_c = self.calc.atoms.get_pbc()
if np.all(pbc_c == False):
self.gamma = True
self.dtype = float
kpts = None
# Multigrid Poisson solver
poisson_solver = PoissonSolver()
else:
if gamma:
self.gamma = True
self.dtype = float
kpts = None
else:
self.gamma = False
self.dtype = complex
# Get k-points from ground-state calculation
kpts = self.calc.input_parameters.kpts
# FFT Poisson solver
poisson_solver = FFTPoissonSolver(dtype=self.dtype)
# K-point descriptor for the q-vectors of the dynamical matrix
# Note, no explicit parallelization here.
self.kd = KPointDescriptor(kpts, 1)
self.kd.set_symmetry(self.atoms, self.symmetry)
self.kd.set_communicator(serial_comm)
# Number of occupied bands
nvalence = self.calc.wfs.nvalence
nbands = nvalence // 2 + nvalence % 2
assert nbands <= self.calc.wfs.bd.nbands
# Extract other useful objects
# Ground-state k-point descriptor - used for the k-points in the
# ResponseCalculator
# XXX replace communicators when ready to parallelize
kd_gs = self.calc.wfs.kd
gd = self.calc.density.gd
kpt_u = self.calc.wfs.kpt_u
setups = self.calc.wfs.setups
dtype_gs = self.calc.wfs.dtype
# WaveFunctions
wfs = WaveFunctions(nbands, kpt_u, setups, kd_gs, gd, dtype=dtype_gs)
# Linear response calculator
self.response_calc = ResponseCalculator(self.calc,
wfs,
dtype=self.dtype)
# Phonon perturbation
self.perturbation = PhononPerturbation(self.calc,
self.kd,
poisson_solver,
dtype=self.dtype)
# Dynamical matrix
self.dyn = DynamicalMatrix(self.atoms, self.kd, dtype=self.dtype)
# Electron-phonon couplings
if e_ph:
self.e_ph = ElectronPhononCoupling(self.atoms,
gd,
self.kd,
dtype=self.dtype)
else:
self.e_ph = None
# Initialization flag
self.initialized = False
# Parallel communicator for parallelization over kpts and domain
self.comm = communicator
def initialize(self):
"""Initialize response calculator and perturbation."""
# Get scaled atomic positions
spos_ac = self.atoms.get_scaled_positions()
self.perturbation.initialize(spos_ac)
self.response_calc.initialize(spos_ac)
self.initialized = True
def __getstate__(self):
"""Method used when pickling.
Bound method attributes cannot be pickled and must therefore be deleted
before an instance is dumped to file.
"""
# Get state of object and take care of troublesome attributes
state = dict(self.__dict__)
state['kd'].__dict__['comm'] = serial_comm
state.pop('calc')
state.pop('perturbation')
state.pop('response_calc')
return state
def run(self, qpts_q=None, clean=False, name=None, path=None):
"""Run calculation for atomic displacements and update matrix.
Parameters
----------
qpts: List
List of q-points indices for which the dynamical matrix will be
calculated (only temporary).
"""
if not self.initialized:
self.initialize()
if self.gamma:
qpts_q = [0]
elif qpts_q is None:
qpts_q = range(self.kd.nibzkpts)
else:
assert isinstance(qpts_q, list)
# Update name and path attributes
self.set_name_and_path(name=name, path=path)
# Get string template for filenames
filename_str = self.get_filename_string()
# Delay the ranks belonging to the same k-point/domain decomposition
# equally
time.sleep(rank // self.comm.size)
# XXX Make a single ground_state_contributions member function
# Ground-state contributions to the force constants
self.dyn.density_ground_state(self.calc)
# self.dyn.wfs_ground_state(self.calc, self.response_calc)
# Calculate linear response wrt q-vectors and displacements of atoms
for q in qpts_q:
if not self.gamma:
self.perturbation.set_q(q)
# First-order contributions to the force constants
for a in self.dyn.indices:
for v in [0, 1, 2]:
# Check if the calculation has already been done
filename = filename_str % (q, a, v)
# Wait for all sub-ranks to enter
self.comm.barrier()
if os.path.isfile(os.path.join(self.path, filename)):
continue
if self.comm.rank == 0:
fd = open(os.path.join(self.path, filename), 'w')
# Wait for all sub-ranks here
self.comm.barrier()
components = ['x', 'y', 'z']
symbols = self.atoms.get_chemical_symbols()
print("q-vector index: %i" % q)
print("Atom index: %i" % a)
print("Atomic symbol: %s" % symbols[a])
print("Component: %s" % components[v])
# Set atom and cartesian component of perturbation
self.perturbation.set_av(a, v)
# Calculate linear response
self.response_calc(self.perturbation)
# Calculate row of the matrix of force constants
self.dyn.calculate_row(self.perturbation,
self.response_calc)
# Write force constants to file
if self.comm.rank == 0:
self.dyn.write(fd, q, a, v)
fd.close()
# Store effective potential derivative
if self.e_ph is not None:
v1_eff_G = self.perturbation.v1_G + \
self.response_calc.vHXC1_G
self.e_ph.v1_eff_qavG.append(v1_eff_G)
# Wait for the file-writing rank here
self.comm.barrier()
# XXX
# Check that all files are valid and collect in a single file
# Remove the files
if clean:
self.clean()
def get_atoms(self):
"""Return atoms."""
return self.atoms
def get_dynamical_matrix(self):
"""Return reference to ``dyn`` attribute."""
return self.dyn
def get_filename_string(self):
"""Return string template for force constant filenames."""
name_str = (self.name + '.' +
'q_%%0%ii_' % len(str(self.kd.nibzkpts)) +
'a_%%0%ii_' % len(str(len(self.atoms))) + 'v_%i' + '.pckl')
return name_str
def set_atoms(self, atoms):
"""Set atoms to be included in the calculation.
Parameters
----------
atoms: list
Can be either a list of strings, ints or ...
"""
assert isinstance(atoms, list)
if isinstance(atoms[0], str):
assert np.all([isinstance(atom, str) for atom in atoms])
sym_a = self.atoms.get_chemical_symbols()
# List for atomic indices
indices = []
for type in atoms:
indices.extend(
[a for a, atom in enumerate(sym_a) if atom == type])
else:
assert np.all([isinstance(atom, int) for atom in atoms])
indices = atoms
self.dyn.set_indices(indices)
def set_name_and_path(self, name=None, path=None):
"""Set name and path of the force constant files.
name: str
Base name for the files which the elements of the matrix of force
constants will be written to.
path: str
Path specifying the directory where the files will be dumped.
"""
if name is None:
self.name = 'phonon.' + self.atoms.get_chemical_formula()
else:
self.name = name
# self.name += '.nibzkpts_%i' % self.kd.nibzkpts
if path is None:
self.path = '.'
else:
self.path = path
# Set corresponding attributes in the ``dyn`` attribute
filename_str = self.get_filename_string()
self.dyn.set_name_and_path(filename_str, self.path)
def clean(self):
"""Delete generated files."""
filename_str = self.get_filename_string()
for q in range(self.kd.nibzkpts):
for a in range(len(self.atoms)):
for v in [0, 1, 2]:
filename = filename_str % (q, a, v)
if os.path.isfile(os.path.join(self.path, filename)):
os.remove(filename)
def band_structure(self, path_kc, modes=False, acoustic=True):
"""Calculate phonon dispersion along a path in the Brillouin zone.
The dynamical matrix at arbitrary q-vectors is obtained by Fourier
transforming the real-space matrix. In case of negative eigenvalues
(squared frequency), the corresponding negative frequency is returned.
Parameters
----------
path_kc: ndarray
List of k-point coordinates (in units of the reciprocal lattice
vectors) specifying the path in the Brillouin zone for which the
dynamical matrix will be calculated.
modes: bool
Returns both frequencies and modes (mass scaled) when True.
acoustic: bool
Restore the acoustic sum-rule in the calculated force constants.
"""
for k_c in path_kc:
assert np.all(np.asarray(k_c) <= 1.0), \
"Scaled coordinates must be given"
# Assemble the dynanical matrix from calculated force constants
self.dyn.assemble(acoustic=acoustic)
# Get the dynamical matrix in real-space
DR_lmn, R_clmn = self.dyn.real_space()
# Reshape for the evaluation of the fourier sums
shape = DR_lmn.shape
DR_m = DR_lmn.reshape((-1, ) + shape[-2:])
R_cm = R_clmn.reshape((3, -1))
# Lists for frequencies and modes along path
omega_kn = []
u_kn = []
# Number of atoms included
N = len(self.dyn.get_indices())
# Mass prefactor for the normal modes
m_inv_av = self.dyn.get_mass_array()
for q_c in path_kc:
# Evaluate fourier transform
phase_m = np.exp(-2.j * pi * np.dot(q_c, R_cm))
# Dynamical matrix in unit of Ha / Bohr**2 / amu
D_q = np.sum(phase_m[:, np.newaxis, np.newaxis] * DR_m, axis=0)
if modes:
omega2_n, u_avn = la.eigh(D_q, UPLO='L')
# Sort eigenmodes according to eigenvalues (see below) and
# multiply with mass prefactor
u_nav = u_avn[:, omega2_n.argsort()].T.copy() * m_inv_av
# Multiply with mass prefactor
u_kn.append(u_nav.reshape((3 * N, -1, 3)))
else:
omega2_n = la.eigvalsh(D_q, UPLO='L')
# Sort eigenvalues in increasing order
omega2_n.sort()
# Use dtype=complex to handle negative eigenvalues
omega_n = np.sqrt(omega2_n.astype(complex))
# Take care of imaginary frequencies
if not np.all(omega2_n >= 0.):
indices = np.where(omega2_n < 0)[0]
print(("WARNING, %i imaginary frequencies at "
"q = (% 5.2f, % 5.2f, % 5.2f) ; (omega_q =% 5.3e*i)" %
(len(indices), q_c[0], q_c[1], q_c[2],
omega_n[indices][0].imag)))
omega_n[indices] = -1 * np.sqrt(np.abs(omega2_n[indices].real))
omega_kn.append(omega_n.real)
# Conversion factor from sqrt(Ha / Bohr**2 / amu) -> eV
s = units.Hartree**0.5 * units._hbar * 1.e10 / \
(units._e * units._amu)**(0.5) / units.Bohr
# Convert to eV and Ang
omega_kn = s * np.asarray(omega_kn)
if modes:
u_kn = np.asarray(u_kn) * units.Bohr
return omega_kn, u_kn
return omega_kn
def write_modes(self,
q_c,
branches=0,
kT=units.kB * 300,
repeat=(1, 1, 1),
nimages=30,
acoustic=True):
"""Write mode to trajectory file.
The classical equipartioning theorem states that each normal mode has
an average energy::
<E> = 1/2 * k_B * T = 1/2 * omega^2 * Q^2
=>
Q = sqrt(k_B*T) / omega
at temperature T. Here, Q denotes the normal coordinate of the mode.
Parameters
----------
q_c: ndarray
q-vector of the modes.
branches: int or list
Branch index of calculated modes.
kT: float
Temperature in units of eV. Determines the amplitude of the atomic
displacements in the modes.
repeat: tuple
Repeat atoms (l, m, n) times in the directions of the lattice
vectors. Displacements of atoms in repeated cells carry a Bloch
phase factor given by the q-vector and the cell lattice vector R_m.
nimages: int
Number of images in an oscillation.
"""
if isinstance(branches, int):
branch_n = [branches]
else:
branch_n = list(branches)
# Calculate modes
omega_n, u_n = self.band_structure([q_c],
modes=True,
acoustic=acoustic)
# Repeat atoms
atoms = self.atoms * repeat
pos_mav = atoms.positions.copy()
# Total number of unit cells
M = np.prod(repeat)
# Corresponding lattice vectors R_m
R_cm = np.indices(repeat[::-1]).reshape(3, -1)[::-1]
# Bloch phase
phase_m = np.exp(2.j * pi * np.dot(q_c, R_cm))
phase_ma = phase_m.repeat(len(self.atoms))
for n in branch_n:
omega = omega_n[0, n]
u_av = u_n[0, n] # .reshape((-1, 3))
# Mean displacement at high T ?
u_av *= sqrt(kT / abs(omega))
mode_av = np.zeros((len(self.atoms), 3), dtype=self.dtype)
indices = self.dyn.get_indices()
mode_av[indices] = u_av
mode_mav = (np.vstack([mode_av] * M) *
phase_ma[:, np.newaxis]).real
traj = Trajectory('%s.mode.%d.traj' % (self.name, n), 'w')
for x in np.linspace(0, 2 * pi, nimages, endpoint=False):
# XXX Is it correct to take out the sine component here ?
atoms.set_positions(pos_mav + sin(x) * mode_mav)
traj.write(atoms)
traj.close()
def create_wave_functions(self, mode, realspace, nspins, nbands, nao,
nvalence, setups, magmom_a, cell_cv, pbc_c):
par = self.parameters
bzkpts_kc = kpts2ndarray(par.kpts, self.atoms)
kd = KPointDescriptor(bzkpts_kc, nspins)
self.timer.start('Set symmetry')
kd.set_symmetry(self.atoms, self.symmetry, comm=self.world)
self.timer.stop('Set symmetry')
self.log(kd)
parallelization = mpi.Parallelization(self.world, nspins * kd.nibzkpts)
parsize_kpt = self.parallel['kpt']
parsize_domain = self.parallel['domain']
parsize_bands = self.parallel['band']
ndomains = None
if parsize_domain is not None:
ndomains = np.prod(parsize_domain)
if mode.name == 'pw':
if ndomains is not None and ndomains > 1:
raise ValueError('Planewave mode does not support '
'domain decomposition.')
ndomains = 1
parallelization.set(kpt=parsize_kpt,
domain=ndomains,
band=parsize_bands)
comms = parallelization.build_communicators()
domain_comm = comms['d']
kpt_comm = comms['k']
band_comm = comms['b']
kptband_comm = comms['D']
domainband_comm = comms['K']
self.comms = comms
if par.gpts is not None:
if par.h is not None:
raise ValueError("""You can't use both "gpts" and "h"!""")
N_c = np.array(par.gpts)
else:
h = par.h
if h is not None:
h /= Bohr
N_c = get_number_of_grid_points(cell_cv, h, mode, realspace,
kd.symmetry)
self.symmetry.check_grid(N_c)
kd.set_communicator(kpt_comm)
parstride_bands = self.parallel['stridebands']
# Unfortunately we need to remember that we adjusted the
# number of bands so we can print a warning if it differs
# from the number specified by the user. (The number can
# be inferred from the input parameters, but it's tricky
# because we allow negative numbers)
self.nbands_parallelization_adjustment = -nbands % band_comm.size
nbands += self.nbands_parallelization_adjustment
bd = BandDescriptor(nbands, band_comm, parstride_bands)
# Construct grid descriptor for coarse grids for wave functions:
gd = self.create_grid_descriptor(N_c, cell_cv, pbc_c, domain_comm,
parsize_domain)
if hasattr(self, 'time') or mode.force_complex_dtype:
dtype = complex
else:
if kd.gamma:
dtype = float
else:
dtype = complex
wfs_kwargs = dict(gd=gd,
nvalence=nvalence,
setups=setups,
bd=bd,
dtype=dtype,
world=self.world,
kd=kd,
kptband_comm=kptband_comm,
timer=self.timer)
if self.parallel['sl_auto']:
# Choose scalapack parallelization automatically
for key, val in self.parallel.items():
if (key.startswith('sl_') and key != 'sl_auto'
and val is not None):
raise ValueError("Cannot use 'sl_auto' together "
"with '%s'" % key)
max_scalapack_cpus = bd.comm.size * gd.comm.size
nprow = max_scalapack_cpus
npcol = 1
# Get a sort of reasonable number of columns/rows
while npcol < nprow and nprow % 2 == 0:
npcol *= 2
nprow //= 2
assert npcol * nprow == max_scalapack_cpus
# ScaLAPACK creates trouble if there aren't at least a few
# whole blocks; choose block size so there will always be
# several blocks. This will crash for small test systems,
# but so will ScaLAPACK in any case
blocksize = min(-(-nbands // 4), 64)
sl_default = (nprow, npcol, blocksize)
else:
sl_default = self.parallel['sl_default']
if mode.name == 'lcao':
# Layouts used for general diagonalizer
sl_lcao = self.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = sl_default
lcaoksl = get_KohnSham_layouts(sl_lcao,
'lcao',
gd,
bd,
domainband_comm,
dtype,
nao=nao,
timer=self.timer)
self.wfs = mode(lcaoksl, **wfs_kwargs)
elif mode.name == 'fd' or mode.name == 'pw':
# buffer_size keyword only relevant for fdpw
buffer_size = self.parallel['buffer_size']
# Layouts used for diagonalizer
sl_diagonalize = self.parallel['sl_diagonalize']
if sl_diagonalize is None:
sl_diagonalize = sl_default
diagksl = get_KohnSham_layouts(
sl_diagonalize,
'fd', # XXX
# choice of key 'fd' not so nice
gd,
bd,
domainband_comm,
dtype,
buffer_size=buffer_size,
timer=self.timer)
# Layouts used for orthonormalizer
sl_inverse_cholesky = self.parallel['sl_inverse_cholesky']
if sl_inverse_cholesky is None:
sl_inverse_cholesky = sl_default
if sl_inverse_cholesky != sl_diagonalize:
message = 'sl_inverse_cholesky != sl_diagonalize ' \
'is not implemented.'
raise NotImplementedError(message)
orthoksl = get_KohnSham_layouts(sl_inverse_cholesky,
'fd',
gd,
bd,
domainband_comm,
dtype,
buffer_size=buffer_size,
timer=self.timer)
# Use (at most) all available LCAO for initialization
lcaonbands = min(nbands, nao // band_comm.size * band_comm.size)
try:
lcaobd = BandDescriptor(lcaonbands, band_comm, parstride_bands)
except RuntimeError:
initksl = None
else:
# Layouts used for general diagonalizer
# (LCAO initialization)
sl_lcao = self.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = sl_default
initksl = get_KohnSham_layouts(sl_lcao,
'lcao',
gd,
lcaobd,
domainband_comm,
dtype,
nao=nao,
timer=self.timer)
self.wfs = mode(diagksl, orthoksl, initksl, **wfs_kwargs)
else:
self.wfs = mode(self, **wfs_kwargs)
self.log(self.wfs, '\n')
def initialize(self, atoms=None):
"""Inexpensive initialization."""
if atoms is None:
atoms = self.atoms
else:
# Save the state of the atoms:
self.atoms = atoms.copy()
par = self.input_parameters
world = par.communicator
if world is None:
world = mpi.world
elif hasattr(world, 'new_communicator'):
# Check for whether object has correct type already
#
# Using isinstance() is complicated because of all the
# combinations, serial/parallel/debug...
pass
else:
# world should be a list of ranks:
world = mpi.world.new_communicator(np.asarray(world))
self.wfs.world = world
self.set_text(par.txt, par.verbose)
natoms = len(atoms)
pos_av = atoms.get_positions() / Bohr
cell_cv = atoms.get_cell()
pbc_c = atoms.get_pbc()
Z_a = atoms.get_atomic_numbers()
magmom_a = atoms.get_initial_magnetic_moments()
magnetic = magmom_a.any()
spinpol = par.spinpol
if par.hund:
if natoms != 1:
raise ValueError('hund=True arg only valid for single atoms!')
spinpol = True
if spinpol is None:
spinpol = magnetic
elif magnetic and not spinpol:
raise ValueError('Non-zero initial magnetic moment for a '
'spin-paired calculation!')
nspins = 1 + int(spinpol)
if isinstance(par.xc, str):
xc = XC(par.xc)
else:
xc = par.xc
setups = Setups(Z_a, par.setups, par.basis, par.lmax, xc, world)
# K-point descriptor
kd = KPointDescriptor(par.kpts, nspins)
width = par.width
if width is None:
if kd.gamma:
width = 0.0
else:
width = 0.1 # eV
else:
assert par.occupations is None
if par.gpts is not None and par.h is None:
N_c = np.array(par.gpts)
else:
if par.h is None:
self.text('Using default value for grid spacing.')
h = 0.2
else:
h = par.h
N_c = h2gpts(h, cell_cv)
cell_cv /= Bohr
if hasattr(self, 'time') or par.dtype==complex:
dtype = complex
else:
if kd.gamma:
dtype = float
else:
dtype = complex
kd.set_symmetry(atoms, setups, par.usesymm, N_c)
nao = setups.nao
nvalence = setups.nvalence - par.charge
nbands = par.nbands
if nbands is None:
nbands = nao
elif nbands > nao and par.mode == 'lcao':
raise ValueError('Too many bands for LCAO calculation: ' +
'%d bands and only %d atomic orbitals!' %
(nbands, nao))
if nvalence < 0:
raise ValueError(
'Charge %f is not possible - not enough valence electrons' %
par.charge)
M = magmom_a.sum()
if par.hund:
f_si = setups[0].calculate_initial_occupation_numbers(
magmom=0, hund=True, charge=par.charge, nspins=nspins)
Mh = f_si[0].sum() - f_si[1].sum()
if magnetic and M != Mh:
raise RuntimeError('You specified a magmom that does not'
'agree with hunds rule!')
else:
M = Mh
if nbands <= 0:
nbands = int(nvalence + M + 0.5) // 2 + (-nbands)
if nvalence > 2 * nbands:
raise ValueError('Too few bands! Electrons: %d, bands: %d'
% (nvalence, nbands))
if par.width is not None:
self.text('**NOTE**: please start using '
'occupations=FermiDirac(width).')
if par.fixmom:
self.text('**NOTE**: please start using '
'occupations=FermiDirac(width, fixmagmom=True).')
if self.occupations is None:
if par.occupations is None:
# Create object for occupation numbers:
self.occupations = occupations.FermiDirac(width, par.fixmom)
else:
self.occupations = par.occupations
self.occupations.magmom = M
cc = par.convergence
if par.mode == 'lcao':
niter_fixdensity = 0
else:
niter_fixdensity = None
if self.scf is None:
self.scf = SCFLoop(
cc['eigenstates'] * nvalence,
cc['energy'] / Hartree * max(nvalence, 1),
cc['density'] * nvalence,
par.maxiter, par.fixdensity,
niter_fixdensity)
parsize, parsize_bands = par.parallel['domain'], par.parallel['band']
if parsize_bands is None:
parsize_bands = 1
# TODO delete/restructure so all checks are in BandDescriptor
if nbands % parsize_bands != 0:
raise RuntimeError('Cannot distribute %d bands to %d processors' %
(nbands, parsize_bands))
if not self.wfs:
if parsize == 'domain only': #XXX this was silly!
parsize = world.size
domain_comm, kpt_comm, band_comm = mpi.distribute_cpus(parsize,
parsize_bands, nspins, kd.nibzkpts, world, par.idiotproof)
kd.set_communicator(kpt_comm)
parstride_bands = par.parallel['stridebands']
bd = BandDescriptor(nbands, band_comm, parstride_bands)
if (self.density is not None and
self.density.gd.comm.size != domain_comm.size):
# Domain decomposition has changed, so we need to
# reinitialize density and hamiltonian:
if par.fixdensity:
raise RuntimeError("I'm confused - please specify parsize."
)
self.density = None
self.hamiltonian = None
# Construct grid descriptor for coarse grids for wave functions:
gd = self.grid_descriptor_class(N_c, cell_cv, pbc_c,
domain_comm, parsize)
# do k-point analysis here? XXX
args = (gd, nvalence, setups, bd, dtype, world, kd, self.timer)
if par.mode == 'lcao':
# Layouts used for general diagonalizer
sl_lcao = par.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = par.parallel['sl_default']
lcaoksl = get_KohnSham_layouts(sl_lcao, 'lcao',
gd, bd, dtype,
nao=nao, timer=self.timer)
self.wfs = LCAOWaveFunctions(lcaoksl, *args)
elif par.mode == 'fd' or isinstance(par.mode, PW):
# buffer_size keyword only relevant for fdpw
buffer_size = par.parallel['buffer_size']
# Layouts used for diagonalizer
sl_diagonalize = par.parallel['sl_diagonalize']
if sl_diagonalize is None:
sl_diagonalize = par.parallel['sl_default']
diagksl = get_KohnSham_layouts(sl_diagonalize, 'fd',
gd, bd, dtype,
buffer_size=buffer_size,
timer=self.timer)
# Layouts used for orthonormalizer
sl_inverse_cholesky = par.parallel['sl_inverse_cholesky']
if sl_inverse_cholesky is None:
sl_inverse_cholesky = par.parallel['sl_default']
if sl_inverse_cholesky != sl_diagonalize:
message = 'sl_inverse_cholesky != sl_diagonalize ' \
'is not implemented.'
raise NotImplementedError(message)
orthoksl = get_KohnSham_layouts(sl_inverse_cholesky, 'fd',
gd, bd, dtype,
buffer_size=buffer_size,
timer=self.timer)
# Use (at most) all available LCAO for initialization
lcaonbands = min(nbands, nao)
lcaobd = BandDescriptor(lcaonbands, band_comm, parstride_bands)
assert nbands <= nao or bd.comm.size == 1
assert lcaobd.mynbands == min(bd.mynbands, nao) #XXX
# Layouts used for general diagonalizer (LCAO initialization)
sl_lcao = par.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = par.parallel['sl_default']
initksl = get_KohnSham_layouts(sl_lcao, 'lcao',
gd, lcaobd, dtype,
nao=nao,
timer=self.timer)
if par.mode == 'fd':
self.wfs = FDWaveFunctions(par.stencils[0], diagksl,
orthoksl, initksl, *args)
else:
# Planewave basis:
self.wfs = par.mode(diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
world, kd, self.timer)
else:
self.wfs = par.mode(self, *args)
else:
self.wfs.set_setups(setups)
if not self.wfs.eigensolver:
# Number of bands to converge:
nbands_converge = cc['bands']
if nbands_converge == 'all':
nbands_converge = nbands
elif nbands_converge != 'occupied':
assert isinstance(nbands_converge, int)
if nbands_converge < 0:
nbands_converge += nbands
eigensolver = get_eigensolver(par.eigensolver, par.mode,
par.convergence)
eigensolver.nbands_converge = nbands_converge
# XXX Eigensolver class doesn't define an nbands_converge property
self.wfs.set_eigensolver(eigensolver)
if self.density is None:
gd = self.wfs.gd
if par.stencils[1] != 9:
# Construct grid descriptor for fine grids for densities
# and potentials:
finegd = gd.refine()
else:
# Special case (use only coarse grid):
finegd = gd
self.density = Density(gd, finegd, nspins,
par.charge + setups.core_charge)
self.density.initialize(setups, par.stencils[1], self.timer,
magmom_a, par.hund)
self.density.set_mixer(par.mixer)
if self.hamiltonian is None:
gd, finegd = self.density.gd, self.density.finegd
self.hamiltonian = Hamiltonian(gd, finegd, nspins,
setups, par.stencils[1], self.timer,
xc, par.poissonsolver,
par.external)
xc.initialize(self.density, self.hamiltonian, self.wfs,
self.occupations)
self.text()
self.print_memory_estimate(self.txt, maxdepth=memory_estimate_depth)
self.txt.flush()
if dry_run:
self.dry_run()
self.initialized = True
class UTKPointParallelSetup(TestCase):
"""
Setup a simple kpoint parallel calculation."""
# Number of bands
nbands = 1
# Spin-polarized
nspins = 1
# Mean spacing and number of grid points per axis (G x G x G)
h = 0.25 / Bohr
G = 48
## Symmetry-reduction of k-points TODO
#symmetry = p.usesymm #XXX 'None' is an allowed value!!!
# Whether spin/k-points are equally distributed (determines nibzkpts)
equipartition = None
nibzkpts = None
gamma = False # can't be gamma point when nibzkpts > 1 ...
dtype = complex #XXX virtual so far..
# =================================
def setUp(self):
for virtvar in ['equipartition']:
assert getattr(self,virtvar) is not None, 'Virtual "%s"!' % virtvar
kpts = {'even' : (12,1,2), \
'prime': (23,1,1)}[self.equipartition]
#primes = [i for i in xrange(50,1,-1) if ~np.any(i%np.arange(2,i)==0)]
bzk_kc = kpts2ndarray(kpts)
assert p.usesymm == None
self.nibzkpts = len(bzk_kc)
#parsize_domain, parsize_bands = create_parsize_minbands(self.nbands, world.size)
parsize_domain, parsize_bands = 1, 1 #XXX
assert self.nbands % np.prod(parsize_bands) == 0
domain_comm, kpt_comm, band_comm = distribute_cpus(parsize_domain,
parsize_bands, self.nspins, self.nibzkpts)
# Set up band descriptor:
self.bd = BandDescriptor(self.nbands, band_comm, p.parallel['stridebands'])
# Set up grid descriptor:
res, ngpts = shapeopt(300, self.G**3, 3, 0.2)
cell_c = self.h * np.array(ngpts)
pbc_c = (True, False, True)
self.gd = GridDescriptor(ngpts, cell_c, pbc_c, domain_comm, parsize_domain)
# Create randomized gas-like atomic configuration
self.atoms = create_random_atoms(self.gd)
# Create setups
Z_a = self.atoms.get_atomic_numbers()
self.setups = Setups(Z_a, p.setups, p.basis, p.lmax, xc)
self.natoms = len(self.setups)
# Set up kpoint descriptor:
self.kd = KPointDescriptor(bzk_kc, self.nspins)
self.kd.set_symmetry(self.atoms, self.setups, usesymm=p.usesymm)
self.kd.set_communicator(kpt_comm)
def tearDown(self):
del self.bd, self.gd, self.kd
del self.setups, self.atoms
def get_parsizes(self): #XXX NO LONGER IN UT_HSOPS?!?
# Careful, overwriting imported GPAW params may cause amnesia in Python.
from gpaw import parsize_domain, parsize_bands
# Choose the largest possible parallelization over kpoint/spins
test_parsize_ks_pairs = gcd(self.nspins*self.nibzkpts, world.size)
remsize = world.size//test_parsize_ks_pairs
# If parsize_bands is not set, choose the largest possible
test_parsize_bands = parsize_bands or gcd(self.nbands, remsize)
# If parsize_bands is not set, choose as few domains as possible
test_parsize_domain = parsize_domain or (remsize//test_parsize_bands)
return test_parsize_domain, test_parsize_bands
# =================================
def verify_comm_sizes(self): #TODO needs work
if world.size == 1:
return
comm_sizes = tuple([comm.size for comm in [world, self.bd.comm, \
self.gd.comm, self.kd.comm]])
self._parinfo = '%d world, %d band, %d domain, %d kpt' % comm_sizes
#self.assertEqual((self.nspins*self.nibzkpts) % self.kd.comm.size, 0) #XXX
def verify_slice_consistency(self):
for kpt_rank in range(self.kd.comm.size):
uslice = self.kd.get_slice(kpt_rank)
myus = np.arange(*uslice.indices(self.kd.nks))
for myu,u in enumerate(myus):
self.assertEqual(self.kd.who_has(u), (kpt_rank, myu))
def verify_combination_consistency(self):
for u in range(self.kd.nks):
s, k = self.kd.what_is(u)
self.assertEqual(self.kd.where_is(s, k), u)
for s in range(self.kd.nspins):
for k in range(self.kd.nibzkpts):
u = self.kd.where_is(s, k)
self.assertEqual(self.kd.what_is(u), (s,k,))
def verify_indexing_consistency(self):
for u in range(self.kd.nks):
kpt_rank, myu = self.kd.who_has(u)
self.assertEqual(self.kd.global_index(myu, kpt_rank), u)
for kpt_rank in range(self.kd.comm.size):
for myu in range(self.kd.get_count(kpt_rank)):
u = self.kd.global_index(myu, kpt_rank)
self.assertEqual(self.kd.who_has(u), (kpt_rank, myu))
def verify_ranking_consistency(self):
ranks = self.kd.get_ranks()
for kpt_rank in range(self.kd.comm.size):
my_indices = self.kd.get_indices(kpt_rank)
matches = np.argwhere(ranks == kpt_rank).ravel()
self.assertTrue((matches == my_indices).all())
for myu in range(self.kd.get_count(kpt_rank)):
u = self.kd.global_index(myu, kpt_rank)
self.assertEqual(my_indices[myu], u)
class UTDomainParallelSetup(TestCase):
"""
Setup a simple domain parallel calculation."""
# Number of bands
nbands = 12
# Spin-polarized
nspins = 1
# Mean spacing and number of grid points per axis (G x G x G)
h = 0.25 / Bohr
G = 48
# Type of boundary conditions employed (determines nibzkpts and dtype)
boundaries = None
nibzkpts = None
dtype = None
timer = nulltimer
def setUp(self):
for virtvar in ['boundaries']:
assert getattr(self,virtvar) is not None, 'Virtual "%s"!' % virtvar
# Basic unit cell information:
res, N_c = shapeopt(100, self.G**3, 3, 0.2)
#N_c = 4*np.round(np.array(N_c)/4) # makes domain decomposition easier
cell_cv = self.h * np.diag(N_c)
pbc_c = {'zero' : (False,False,False), \
'periodic': (True,True,True), \
'mixed' : (True, False, True)}[self.boundaries]
# Create randomized gas-like atomic configuration on interim grid
tmpgd = GridDescriptor(N_c, cell_cv, pbc_c)
self.atoms = create_random_atoms(tmpgd)
# Create setups
Z_a = self.atoms.get_atomic_numbers()
assert 1 == self.nspins
self.setups = Setups(Z_a, p.setups, p.basis, p.lmax, xc)
self.natoms = len(self.setups)
# Decide how many kpoints to sample from the 1st Brillouin Zone
kpts_c = np.ceil((10/Bohr)/np.sum(cell_cv**2,axis=1)**0.5).astype(int)
kpts_c = tuple(kpts_c * pbc_c + 1 - pbc_c)
self.bzk_kc = kpts2ndarray(kpts_c)
# Set up k-point descriptor
self.kd = KPointDescriptor(self.bzk_kc, self.nspins)
self.kd.set_symmetry(self.atoms, self.setups, p.usesymm)
# Set the dtype
if self.kd.gamma:
self.dtype = float
else:
self.dtype = complex
# Create communicators
parsize, parsize_bands = self.get_parsizes()
assert self.nbands % np.prod(parsize_bands) == 0
domain_comm, kpt_comm, band_comm = distribute_cpus(parsize,
parsize_bands, self.nspins, self.kd.nibzkpts)
self.kd.set_communicator(kpt_comm)
# Set up band descriptor:
self.bd = BandDescriptor(self.nbands, band_comm)
# Set up grid descriptor:
self.gd = GridDescriptor(N_c, cell_cv, pbc_c, domain_comm, parsize)
# Set up kpoint/spin descriptor (to be removed):
self.kd_old = KPointDescriptorOld(self.nspins, self.kd.nibzkpts,
kpt_comm, self.kd.gamma, self.dtype)
def tearDown(self):
del self.atoms, self.bd, self.gd, self.kd, self.kd_old
def get_parsizes(self):
# Careful, overwriting imported GPAW params may cause amnesia in Python.
from gpaw import parsize, parsize_bands
# D: number of domains
# B: number of band groups
if parsize is None:
D = min(world.size, 2)
else:
D = parsize
assert world.size % D == 0
if parsize_bands is None:
B = world.size // D
else:
B = parsize_bands
return D, B
# =================================
def verify_comm_sizes(self):
if world.size == 1:
return
comm_sizes = tuple([comm.size for comm in [world, self.bd.comm, \
self.gd.comm, self.kd_old.comm]])
self._parinfo = '%d world, %d band, %d domain, %d kpt' % comm_sizes
self.assertEqual(self.nbands % self.bd.comm.size, 0)
self.assertEqual((self.nspins * self.kd.nibzkpts) % self.kd_old.comm.size, 0)
class UTKPointParallelSetup(TestCase):
"""
Setup a simple kpoint parallel calculation."""
# Number of bands
nbands = 1
# Spin-polarized
nspins = 1
# Mean spacing and number of grid points per axis (G x G x G)
h = 0.25 / Bohr
G = 48
## Symmetry-reduction of k-points TODO
#symmetry = p.usesymm #XXX 'None' is an allowed value!!!
# Whether spin/k-points are equally distributed (determines nibzkpts)
equipartition = None
nibzkpts = None
gamma = False # can't be gamma point when nibzkpts > 1 ...
dtype = complex #XXX virtual so far..
# =================================
def setUp(self):
for virtvar in ['equipartition']:
assert getattr(self,virtvar) is not None, 'Virtual "%s"!' % virtvar
kpts = {'even' : (12,1,2), \
'prime': (23,1,1)}[self.equipartition]
#primes = [i for i in xrange(50,1,-1) if ~np.any(i%np.arange(2,i)==0)]
bzk_kc = kpts2ndarray(kpts)
assert p.usesymm == None
self.nibzkpts = len(bzk_kc)
#parsize, parsize_bands = create_parsize_minbands(self.nbands, world.size)
parsize, parsize_bands = 1, 1 #XXX
assert self.nbands % np.prod(parsize_bands) == 0
domain_comm, kpt_comm, band_comm = distribute_cpus(parsize,
parsize_bands, self.nspins, self.nibzkpts)
# Set up band descriptor:
self.bd = BandDescriptor(self.nbands, band_comm, p.parallel['stridebands'])
# Set up grid descriptor:
res, ngpts = shapeopt(300, self.G**3, 3, 0.2)
cell_c = self.h * np.array(ngpts)
pbc_c = (True, False, True)
self.gd = GridDescriptor(ngpts, cell_c, pbc_c, domain_comm, parsize)
# Create randomized gas-like atomic configuration
self.atoms = create_random_atoms(self.gd)
# Create setups
Z_a = self.atoms.get_atomic_numbers()
self.setups = Setups(Z_a, p.setups, p.basis, p.lmax, xc)
self.natoms = len(self.setups)
# Set up kpoint descriptor:
self.kd = KPointDescriptor(bzk_kc, self.nspins)
self.kd.set_symmetry(self.atoms, self.setups, p.usesymm)
self.kd.set_communicator(kpt_comm)
def tearDown(self):
del self.bd, self.gd, self.kd
del self.setups, self.atoms
def get_parsizes(self): #XXX NO LONGER IN UT_HSOPS?!?
# Careful, overwriting imported GPAW params may cause amnesia in Python.
from gpaw import parsize, parsize_bands
# Choose the largest possible parallelization over kpoint/spins
test_parsize_ks_pairs = gcd(self.nspins*self.nibzkpts, world.size)
remsize = world.size//test_parsize_ks_pairs
# If parsize_bands is not set, choose the largest possible
test_parsize_bands = parsize_bands or gcd(self.nbands, remsize)
# If parsize_bands is not set, choose as few domains as possible
test_parsize = parsize or (remsize//test_parsize_bands)
return test_parsize, test_parsize_bands
# =================================
def verify_comm_sizes(self): #TODO needs work
if world.size == 1:
return
comm_sizes = tuple([comm.size for comm in [world, self.bd.comm, \
self.gd.comm, self.kd.comm]])
self._parinfo = '%d world, %d band, %d domain, %d kpt' % comm_sizes
#self.assertEqual((self.nspins*self.nibzkpts) % self.kd.comm.size, 0) #XXX
def verify_slice_consistency(self):
for kpt_rank in range(self.kd.comm.size):
uslice = self.kd.get_slice(kpt_rank)
myus = np.arange(*uslice.indices(self.kd.nks))
for myu,u in enumerate(myus):
self.assertEqual(self.kd.who_has(u), (kpt_rank, myu))
def verify_combination_consistency(self):
for u in range(self.kd.nks):
s, k = self.kd.what_is(u)
self.assertEqual(self.kd.where_is(s, k), u)
for s in range(self.kd.nspins):
for k in range(self.kd.nibzkpts):
u = self.kd.where_is(s, k)
self.assertEqual(self.kd.what_is(u), (s,k,))
def verify_indexing_consistency(self):
for u in range(self.kd.nks):
kpt_rank, myu = self.kd.who_has(u)
self.assertEqual(self.kd.global_index(myu, kpt_rank), u)
for kpt_rank in range(self.kd.comm.size):
for myu in range(self.kd.get_count(kpt_rank)):
u = self.kd.global_index(myu, kpt_rank)
self.assertEqual(self.kd.who_has(u), (kpt_rank, myu))
def verify_ranking_consistency(self):
ranks = self.kd.get_ranks()
for kpt_rank in range(self.kd.comm.size):
my_indices = self.kd.get_indices(kpt_rank)
matches = np.argwhere(ranks == kpt_rank).ravel()
self.assertTrue((matches == my_indices).all())
for myu in range(self.kd.get_count(kpt_rank)):
u = self.kd.global_index(myu, kpt_rank)
self.assertEqual(my_indices[myu], u)
def initialize(self, atoms=None):
"""Inexpensive initialization."""
if atoms is None:
atoms = self.atoms
else:
# Save the state of the atoms:
self.atoms = atoms.copy()
par = self.input_parameters
world = par.communicator
if world is None:
world = mpi.world
elif hasattr(world, 'new_communicator'):
# Check for whether object has correct type already
#
# Using isinstance() is complicated because of all the
# combinations, serial/parallel/debug...
pass
else:
# world should be a list of ranks:
world = mpi.world.new_communicator(np.asarray(world))
self.wfs.world = world
if 'txt' in self._changed_keywords:
self.set_txt(par.txt)
self.verbose = par.verbose
natoms = len(atoms)
cell_cv = atoms.get_cell() / Bohr
pbc_c = atoms.get_pbc()
Z_a = atoms.get_atomic_numbers()
magmom_av = atoms.get_initial_magnetic_moments()
self.check_atoms()
# Generate new xc functional only when it is reset by set
# XXX sounds like this should use the _changed_keywords dictionary.
if self.hamiltonian is None or self.hamiltonian.xc is None:
if isinstance(par.xc, str):
xc = XC(par.xc)
else:
xc = par.xc
else:
xc = self.hamiltonian.xc
mode = par.mode
if mode == 'fd':
mode = FD()
elif mode == 'pw':
mode = pw.PW()
elif mode == 'lcao':
mode = LCAO()
else:
assert hasattr(mode, 'name'), str(mode)
if xc.orbital_dependent and mode.name == 'lcao':
raise NotImplementedError('LCAO mode does not support '
'orbital-dependent XC functionals.')
if par.realspace is None:
realspace = (mode.name != 'pw')
else:
realspace = par.realspace
if mode.name == 'pw':
assert not realspace
if par.filter is None and mode.name != 'pw':
gamma = 1.6
if par.gpts is not None:
h = ((np.linalg.inv(cell_cv)**2).sum(0)**-0.5 / par.gpts).max()
else:
h = (par.h or 0.2) / Bohr
def filter(rgd, rcut, f_r, l=0):
gcut = np.pi / h - 2 / rcut / gamma
f_r[:] = rgd.filter(f_r, rcut * gamma, gcut, l)
else:
filter = par.filter
setups = Setups(Z_a, par.setups, par.basis, par.lmax, xc, filter,
world)
if magmom_av.ndim == 1:
collinear = True
magmom_av, magmom_a = np.zeros((natoms, 3)), magmom_av
magmom_av[:, 2] = magmom_a
else:
collinear = False
magnetic = magmom_av.any()
spinpol = par.spinpol
if par.hund:
if natoms != 1:
raise ValueError('hund=True arg only valid for single atoms!')
spinpol = True
magmom_av[0] = (0, 0, setups[0].get_hunds_rule_moment(par.charge))
if spinpol is None:
spinpol = magnetic
elif magnetic and not spinpol:
raise ValueError('Non-zero initial magnetic moment for a ' +
'spin-paired calculation!')
if collinear:
nspins = 1 + int(spinpol)
ncomp = 1
else:
nspins = 1
ncomp = 2
if par.usesymm != 'default':
warnings.warn('Use "symmetry" keyword instead of ' +
'"usesymm" keyword')
par.symmetry = usesymm2symmetry(par.usesymm)
symm = par.symmetry
if symm == 'off':
symm = {'point_group': False, 'time_reversal': False}
bzkpts_kc = kpts2ndarray(par.kpts, self.atoms)
kd = KPointDescriptor(bzkpts_kc, nspins, collinear)
m_av = magmom_av.round(decimals=3) # round off
id_a = zip(setups.id_a, *m_av.T)
symmetry = Symmetry(id_a, cell_cv, atoms.pbc, **symm)
kd.set_symmetry(atoms, symmetry, comm=world)
setups.set_symmetry(symmetry)
if par.gpts is not None:
N_c = np.array(par.gpts)
else:
h = par.h
if h is not None:
h /= Bohr
N_c = get_number_of_grid_points(cell_cv, h, mode, realspace,
kd.symmetry)
symmetry.check_grid(N_c)
width = par.width
if width is None:
if pbc_c.any():
width = 0.1 # eV
else:
width = 0.0
else:
assert par.occupations is None
if hasattr(self, 'time') or par.dtype == complex:
dtype = complex
else:
if kd.gamma:
dtype = float
else:
dtype = complex
nao = setups.nao
nvalence = setups.nvalence - par.charge
M_v = magmom_av.sum(0)
M = np.dot(M_v, M_v)**0.5
nbands = par.nbands
orbital_free = any(setup.orbital_free for setup in setups)
if orbital_free:
nbands = 1
if isinstance(nbands, basestring):
if nbands[-1] == '%':
basebands = int(nvalence + M + 0.5) // 2
nbands = int((float(nbands[:-1]) / 100) * basebands)
else:
raise ValueError('Integer Expected: Only use a string '
'if giving a percentage of occupied bands')
if nbands is None:
nbands = 0
for setup in setups:
nbands_from_atom = setup.get_default_nbands()
# Any obscure setup errors?
if nbands_from_atom < -(-setup.Nv // 2):
raise ValueError('Bad setup: This setup requests %d'
' bands but has %d electrons.' %
(nbands_from_atom, setup.Nv))
nbands += nbands_from_atom
nbands = min(nao, nbands)
elif nbands > nao and mode.name == 'lcao':
raise ValueError('Too many bands for LCAO calculation: '
'%d bands and only %d atomic orbitals!' %
(nbands, nao))
if nvalence < 0:
raise ValueError(
'Charge %f is not possible - not enough valence electrons' %
par.charge)
if nbands <= 0:
nbands = int(nvalence + M + 0.5) // 2 + (-nbands)
if nvalence > 2 * nbands and not orbital_free:
raise ValueError('Too few bands! Electrons: %f, bands: %d' %
(nvalence, nbands))
nbands *= ncomp
if par.width is not None:
self.text('**NOTE**: please start using '
'occupations=FermiDirac(width).')
if par.fixmom:
self.text('**NOTE**: please start using '
'occupations=FermiDirac(width, fixmagmom=True).')
if self.occupations is None:
if par.occupations is None:
# Create object for occupation numbers:
if orbital_free:
width = 0.0 # even for PBC
self.occupations = occupations.TFOccupations(
width, par.fixmom)
else:
self.occupations = occupations.FermiDirac(
width, par.fixmom)
else:
self.occupations = par.occupations
# If occupation numbers are changed, and we have wave functions,
# recalculate the occupation numbers
if self.wfs is not None and not isinstance(self.wfs,
EmptyWaveFunctions):
self.occupations.calculate(self.wfs)
self.occupations.magmom = M_v[2]
cc = par.convergence
if mode.name == 'lcao':
niter_fixdensity = 0
else:
niter_fixdensity = None
if self.scf is None:
force_crit = cc['forces']
if force_crit is not None:
force_crit /= Hartree / Bohr
self.scf = SCFLoop(cc['eigenstates'] / Hartree**2 * nvalence,
cc['energy'] / Hartree * max(nvalence, 1),
cc['density'] * nvalence, par.maxiter,
par.fixdensity, niter_fixdensity, force_crit)
parsize_kpt = par.parallel['kpt']
parsize_domain = par.parallel['domain']
parsize_bands = par.parallel['band']
if not realspace:
pbc_c = np.ones(3, bool)
if not self.wfs:
if parsize_domain == 'domain only': # XXX this was silly!
parsize_domain = world.size
parallelization = mpi.Parallelization(world, nspins * kd.nibzkpts)
ndomains = None
if parsize_domain is not None:
ndomains = np.prod(parsize_domain)
if mode.name == 'pw':
if ndomains > 1:
raise ValueError('Planewave mode does not support '
'domain decomposition.')
ndomains = 1
parallelization.set(kpt=parsize_kpt,
domain=ndomains,
band=parsize_bands)
comms = parallelization.build_communicators()
domain_comm = comms['d']
kpt_comm = comms['k']
band_comm = comms['b']
kptband_comm = comms['D']
domainband_comm = comms['K']
self.comms = comms
kd.set_communicator(kpt_comm)
parstride_bands = par.parallel['stridebands']
# Unfortunately we need to remember that we adjusted the
# number of bands so we can print a warning if it differs
# from the number specified by the user. (The number can
# be inferred from the input parameters, but it's tricky
# because we allow negative numbers)
self.nbands_parallelization_adjustment = -nbands % band_comm.size
nbands += self.nbands_parallelization_adjustment
# I would like to give the following error message, but apparently
# there are cases, e.g. gpaw/test/gw_ppa.py, which involve
# nbands > nao and are supposed to work that way.
#if nbands > nao:
# raise ValueError('Number of bands %d adjusted for band '
# 'parallelization %d exceeds number of atomic '
# 'orbitals %d. This problem can be fixed '
# 'by reducing the number of bands a bit.'
# % (nbands, band_comm.size, nao))
bd = BandDescriptor(nbands, band_comm, parstride_bands)
if (self.density is not None
and self.density.gd.comm.size != domain_comm.size):
# Domain decomposition has changed, so we need to
# reinitialize density and hamiltonian:
if par.fixdensity:
raise RuntimeError(
'Density reinitialization conflict ' +
'with "fixdensity" - specify domain decomposition.')
self.density = None
self.hamiltonian = None
# Construct grid descriptor for coarse grids for wave functions:
gd = self.grid_descriptor_class(N_c, cell_cv, pbc_c, domain_comm,
parsize_domain)
# do k-point analysis here? XXX
args = (gd, nvalence, setups, bd, dtype, world, kd, kptband_comm,
self.timer)
if par.parallel['sl_auto']:
# Choose scalapack parallelization automatically
for key, val in par.parallel.items():
if (key.startswith('sl_') and key != 'sl_auto'
and val is not None):
raise ValueError("Cannot use 'sl_auto' together "
"with '%s'" % key)
max_scalapack_cpus = bd.comm.size * gd.comm.size
nprow = max_scalapack_cpus
npcol = 1
# Get a sort of reasonable number of columns/rows
while npcol < nprow and nprow % 2 == 0:
npcol *= 2
nprow //= 2
assert npcol * nprow == max_scalapack_cpus
# ScaLAPACK creates trouble if there aren't at least a few
# whole blocks; choose block size so there will always be
# several blocks. This will crash for small test systems,
# but so will ScaLAPACK in any case
blocksize = min(-(-nbands // 4), 64)
sl_default = (nprow, npcol, blocksize)
else:
sl_default = par.parallel['sl_default']
if mode.name == 'lcao':
# Layouts used for general diagonalizer
sl_lcao = par.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = sl_default
lcaoksl = get_KohnSham_layouts(sl_lcao,
'lcao',
gd,
bd,
domainband_comm,
dtype,
nao=nao,
timer=self.timer)
self.wfs = mode(collinear, lcaoksl, *args)
elif mode.name == 'fd' or mode.name == 'pw':
# buffer_size keyword only relevant for fdpw
buffer_size = par.parallel['buffer_size']
# Layouts used for diagonalizer
sl_diagonalize = par.parallel['sl_diagonalize']
if sl_diagonalize is None:
sl_diagonalize = sl_default
diagksl = get_KohnSham_layouts(
sl_diagonalize,
'fd', # XXX
# choice of key 'fd' not so nice
gd,
bd,
domainband_comm,
dtype,
buffer_size=buffer_size,
timer=self.timer)
# Layouts used for orthonormalizer
sl_inverse_cholesky = par.parallel['sl_inverse_cholesky']
if sl_inverse_cholesky is None:
sl_inverse_cholesky = sl_default
if sl_inverse_cholesky != sl_diagonalize:
message = 'sl_inverse_cholesky != sl_diagonalize ' \
'is not implemented.'
raise NotImplementedError(message)
orthoksl = get_KohnSham_layouts(sl_inverse_cholesky,
'fd',
gd,
bd,
domainband_comm,
dtype,
buffer_size=buffer_size,
timer=self.timer)
# Use (at most) all available LCAO for initialization
lcaonbands = min(nbands, nao)
try:
lcaobd = BandDescriptor(lcaonbands, band_comm,
parstride_bands)
except RuntimeError:
initksl = None
else:
# Layouts used for general diagonalizer
# (LCAO initialization)
sl_lcao = par.parallel['sl_lcao']
if sl_lcao is None:
sl_lcao = sl_default
initksl = get_KohnSham_layouts(sl_lcao,
'lcao',
gd,
lcaobd,
domainband_comm,
dtype,
nao=nao,
timer=self.timer)
if hasattr(self, 'time'):
assert mode.name == 'fd'
from gpaw.tddft import TimeDependentWaveFunctions
self.wfs = TimeDependentWaveFunctions(
par.stencils[0], diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, world, kd, kptband_comm,
self.timer)
elif mode.name == 'fd':
self.wfs = mode(par.stencils[0], diagksl, orthoksl,
initksl, *args)
else:
assert mode.name == 'pw'
self.wfs = mode(diagksl, orthoksl, initksl, *args)
else:
self.wfs = mode(self, *args)
else:
self.wfs.set_setups(setups)
if not self.wfs.eigensolver:
# Number of bands to converge:
nbands_converge = cc['bands']
if nbands_converge == 'all':
nbands_converge = nbands
elif nbands_converge != 'occupied':
assert isinstance(nbands_converge, int)
if nbands_converge < 0:
nbands_converge += nbands
eigensolver = get_eigensolver(par.eigensolver, mode,
par.convergence)
eigensolver.nbands_converge = nbands_converge
# XXX Eigensolver class doesn't define an nbands_converge property
if isinstance(xc, SIC):
eigensolver.blocksize = 1
self.wfs.set_eigensolver(eigensolver)
if self.density is None:
gd = self.wfs.gd
if par.stencils[1] != 9:
# Construct grid descriptor for fine grids for densities
# and potentials:
finegd = gd.refine()
else:
# Special case (use only coarse grid):
finegd = gd
if realspace:
self.density = RealSpaceDensity(
gd, finegd, nspins, par.charge + setups.core_charge,
collinear, par.stencils[1])
else:
self.density = pw.ReciprocalSpaceDensity(
gd, finegd, nspins, par.charge + setups.core_charge,
collinear)
self.density.initialize(setups, self.timer, magmom_av, par.hund)
self.density.set_mixer(par.mixer)
if self.hamiltonian is None:
gd, finegd = self.density.gd, self.density.finegd
if realspace:
self.hamiltonian = RealSpaceHamiltonian(
gd, finegd, nspins, setups, self.timer, xc, world,
self.wfs.kptband_comm, par.external, collinear,
par.poissonsolver, par.stencils[1])
else:
self.hamiltonian = pw.ReciprocalSpaceHamiltonian(
gd, finegd, self.density.pd2, self.density.pd3, nspins,
setups, self.timer, xc, world, self.wfs.kptband_comm,
par.external, collinear)
xc.initialize(self.density, self.hamiltonian, self.wfs,
self.occupations)
self.text()
self.print_memory_estimate(self.txt, maxdepth=memory_estimate_depth)
self.txt.flush()
self.timer.print_info(self)
if dry_run:
self.dry_run()
if realspace and \
self.hamiltonian.poisson.get_description() == 'FDTD+TDDFT':
self.hamiltonian.poisson.set_density(self.density)
self.hamiltonian.poisson.print_messages(self.text)
self.txt.flush()
self.initialized = True
self._changed_keywords.clear()
class HybridXC(XCFunctional):
orbital_dependent = True
def __init__(self, name, hybrid=None, xc=None, finegrid=False, alpha=None):
"""Mix standard functionals with exact exchange.
name: str
Name of hybrid functional.
hybrid: float
Fraction of exact exchange.
xc: str or XCFunctional object
Standard DFT functional with scaled down exchange.
finegrid: boolean
Use fine grid for energy functional evaluations?
"""
if name == 'EXX':
assert hybrid is None and xc is None
hybrid = 1.0
xc = XC(XCNull())
elif name == 'PBE0':
assert hybrid is None and xc is None
hybrid = 0.25
xc = XC('HYB_GGA_XC_PBEH')
elif name == 'B3LYP':
assert hybrid is None and xc is None
hybrid = 0.2
xc = XC('HYB_GGA_XC_B3LYP')
if isinstance(xc, str):
xc = XC(xc)
self.hybrid = hybrid
self.xc = xc
self.type = xc.type
self.alpha = alpha
self.exx = 0.0
XCFunctional.__init__(self, name)
def get_setup_name(self):
return 'PBE'
def calculate_radial(self,
rgd,
n_sLg,
Y_L,
v_sg,
dndr_sLg=None,
rnablaY_Lv=None,
tau_sg=None,
dedtau_sg=None):
return self.xc.calculate_radial(rgd, n_sLg, Y_L, v_sg, dndr_sLg,
rnablaY_Lv)
def initialize(self, density, hamiltonian, wfs, occupations):
self.xc.initialize(density, hamiltonian, wfs, occupations)
self.nspins = wfs.nspins
self.setups = wfs.setups
self.density = density
self.kpt_u = wfs.kpt_u
self.gd = density.gd
self.kd = wfs.kd
self.bd = wfs.bd
N_c = self.gd.N_c
N = self.gd.N_c.prod()
vol = self.gd.dv * N
if self.alpha is None:
self.alpha = 6 * vol**(2 / 3.0) / pi**2
self.gamma = (vol / (2 * pi)**2 * sqrt(pi / self.alpha) *
self.kd.nbzkpts)
ecut = 0.5 * pi**2 / (self.gd.h_cv**2).sum(1).max()
if self.kd.N_c is None:
self.bzk_kc = np.zeros((1, 3))
dfghdfgh
else:
n = self.kd.N_c * 2 - 1
bzk_kc = np.indices(n).transpose((1, 2, 3, 0))
bzk_kc.shape = (-1, 3)
bzk_kc -= self.kd.N_c - 1
self.bzk_kc = bzk_kc.astype(float) / self.kd.N_c
self.pwd = PWDescriptor(ecut, self.gd, self.bzk_kc)
n = 0
for k_c, Gpk2_G in zip(self.bzk_kc[:], self.pwd.G2_qG):
if (k_c > -0.5).all() and (k_c <= 0.5).all(): #XXX???
if k_c.any():
self.gamma -= np.dot(np.exp(-self.alpha * Gpk2_G),
Gpk2_G**-1)
else:
self.gamma -= np.dot(np.exp(-self.alpha * Gpk2_G[1:]),
Gpk2_G[1:]**-1)
n += 1
assert n == self.kd.N_c.prod()
self.ghat = LFC(self.gd, [setup.ghat_l for setup in density.setups],
dtype=complex)
self.ghat.set_k_points(self.bzk_kc)
self.fullkd = KPointDescriptor(self.kd.bzk_kc, nspins=1)
class S:
id_a = []
def set_symmetry(self, s):
pass
self.fullkd.set_symmetry(Atoms(pbc=True), S(), False)
self.fullkd.set_communicator(world)
self.pt = LFC(self.gd, [setup.pt_j for setup in density.setups],
dtype=complex)
self.pt.set_k_points(self.fullkd.ibzk_kc)
self.interpolator = density.interpolator
def set_positions(self, spos_ac):
self.ghat.set_positions(spos_ac)
self.pt.set_positions(spos_ac)
def calculate(self, gd, n_sg, v_sg=None, e_g=None):
# Normal XC contribution:
exc = self.xc.calculate(gd, n_sg, v_sg, e_g)
# Add EXX contribution:
return exc + self.exx
def calculate_exx(self):
"""Non-selfconsistent calculation."""
kd = self.kd
K = self.fullkd.nibzkpts
assert self.nspins == 1
Q = K // world.size
assert Q * world.size == K
parallel = (world.size > self.nspins)
self.exx = 0.0
self.exx_skn = np.zeros((self.nspins, K, self.bd.nbands))
kpt_u = []
for k in range(world.rank * Q, (world.rank + 1) * Q):
k_c = self.fullkd.ibzk_kc[k]
for k1, k1_c in enumerate(kd.bzk_kc):
if abs(k1_c - k_c).max() < 1e-10:
break
# Index of symmetry related point in the irreducible BZ
ik = kd.kibz_k[k1]
kpt = self.kpt_u[ik]
# KPoint from ground-state calculation
phase_cd = np.exp(2j * pi * self.gd.sdisp_cd * k_c[:, np.newaxis])
kpt2 = KPoint0(kpt.weight, kpt.s, k, None, phase_cd)
kpt2.psit_nG = np.empty_like(kpt.psit_nG)
kpt2.f_n = kpt.f_n / kpt.weight / K * 2
for n, psit_G in enumerate(kpt2.psit_nG):
psit_G[:] = kd.transform_wave_function(kpt.psit_nG[n], k1)
kpt2.P_ani = self.pt.dict(len(kpt.psit_nG))
self.pt.integrate(kpt2.psit_nG, kpt2.P_ani, k)
kpt_u.append(kpt2)
for s in range(self.nspins):
kpt1_q = [KPoint(self.fullkd, kpt) for kpt in kpt_u if kpt.s == s]
kpt2_q = kpt1_q[:]
if len(kpt1_q) == 0:
# No s-spins on this CPU:
continue
# Send rank:
srank = self.fullkd.get_rank_and_index(s, (kpt1_q[0].k - 1) % K)[0]
# Receive rank:
rrank = self.fullkd.get_rank_and_index(s,
(kpt1_q[-1].k + 1) % K)[0]
# Shift k-points K // 2 times:
for i in range(K // 2 + 1):
if i < K // 2:
if parallel:
kpt = kpt2_q[-1].next()
kpt.start_receiving(rrank)
kpt2_q[0].start_sending(srank)
else:
kpt = kpt2_q[0]
for kpt1, kpt2 in zip(kpt1_q, kpt2_q):
if 2 * i == K:
self.apply(kpt1, kpt2, invert=(kpt1.k > kpt2.k))
else:
self.apply(kpt1, kpt2)
self.apply(kpt1, kpt2, invert=True)
if i < K // 2:
if parallel:
kpt.wait()
kpt2_q[0].wait()
kpt2_q.pop(0)
kpt2_q.append(kpt)
self.exx = world.sum(self.exx)
world.sum(self.exx_skn)
self.exx += self.calculate_paw_correction()
def apply(self, kpt1, kpt2, invert=False):
#print world.rank,kpt1.k,kpt2.k,invert
k1_c = self.fullkd.ibzk_kc[kpt1.k]
k2_c = self.fullkd.ibzk_kc[kpt2.k]
if invert:
k2_c = -k2_c
k12_c = k1_c - k2_c
N_c = self.gd.N_c
eikr_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k12_c / N_c).T)
for q, k_c in enumerate(self.bzk_kc):
if abs(k_c + k12_c).max() < 1e-9:
q0 = q
break
for q, k_c in enumerate(self.bzk_kc):
if abs(k_c - k12_c).max() < 1e-9:
q00 = q
break
Gpk2_G = self.pwd.G2_qG[q0]
if Gpk2_G[0] == 0:
Gpk2_G = Gpk2_G.copy()
Gpk2_G[0] = 1.0 / self.gamma
N = N_c.prod()
vol = self.gd.dv * N
nspins = self.nspins
same = (kpt1.k == kpt2.k)
for n1, psit1_R in enumerate(kpt1.psit_nG):
f1 = kpt1.f_n[n1]
for n2, psit2_R in enumerate(kpt2.psit_nG):
if same and n2 > n1:
continue
f2 = kpt2.f_n[n2]
nt_R = self.calculate_pair_density(n1, n2, kpt1, kpt2, q0,
invert)
nt_G = self.pwd.fft(nt_R * eikr_R) / N
vt_G = nt_G.copy()
vt_G *= -pi * vol / Gpk2_G
e = np.vdot(nt_G, vt_G).real * nspins * self.hybrid
if same and n1 == n2:
e /= 2
self.exx += e * f1 * f2
self.ekin -= 2 * e * f1 * f2
self.exx_skn[kpt1.s, kpt1.k, n1] += f2 * e
self.exx_skn[kpt2.s, kpt2.k, n2] += f1 * e
calculate_potential = not True
if calculate_potential:
vt_R = self.pwd.ifft(vt_G).conj() * eikr_R * N / vol
if kpt1 is kpt2 and not invert and n1 == n2:
kpt1.vt_nG[n1] = 0.5 * f1 * vt_R
if invert:
kpt1.Htpsit_nG[n1] += \
f2 * nspins * psit2_R.conj() * vt_R
else:
kpt1.Htpsit_nG[n1] += f2 * nspins * psit2_R * vt_R
if kpt1 is not kpt2:
if invert:
kpt2.Htpsit_nG[n2] += (f1 * nspins *
psit1_R.conj() * vt_R)
else:
kpt2.Htpsit_nG[n2] += (f1 * nspins * psit1_R *
vt_R.conj())
def calculate_paw_correction(self):
exx = 0
deg = 2 // self.nspins # spin degeneracy
for a, D_sp in self.density.D_asp.items():
setup = self.setups[a]
for D_p in D_sp:
D_ii = unpack2(D_p)
ni = len(D_ii)
for i1 in range(ni):
for i2 in range(ni):
A = 0.0
for i3 in range(ni):
p13 = packed_index(i1, i3, ni)
for i4 in range(ni):
p24 = packed_index(i2, i4, ni)
A += setup.M_pp[p13, p24] * D_ii[i3, i4]
p12 = packed_index(i1, i2, ni)
exx -= self.hybrid / deg * D_ii[i1, i2] * A
if setup.X_p is not None:
exx -= self.hybrid * np.dot(D_p, setup.X_p)
exx += self.hybrid * setup.ExxC
return exx
def calculate_pair_density(self, n1, n2, kpt1, kpt2, q, invert):
if invert:
nt_G = kpt1.psit_nG[n1].conj() * kpt2.psit_nG[n2].conj()
else:
nt_G = kpt1.psit_nG[n1].conj() * kpt2.psit_nG[n2]
Q_aL = {}
for a, P1_ni in kpt1.P_ani.items():
P1_i = P1_ni[n1]
P2_i = kpt2.P_ani[a][n2]
if invert:
D_ii = np.outer(P1_i.conj(), P2_i.conj())
else:
D_ii = np.outer(P1_i.conj(), P2_i)
D_p = pack(D_ii)
Q_aL[a] = np.dot(D_p, self.setups[a].Delta_pL)
self.ghat.add(nt_G, Q_aL, q)
return nt_G
class HybridXC(XCFunctional):
orbital_dependent = True
def __init__(self, name, hybrid=None, xc=None, finegrid=False,
alpha=None):
"""Mix standard functionals with exact exchange.
name: str
Name of hybrid functional.
hybrid: float
Fraction of exact exchange.
xc: str or XCFunctional object
Standard DFT functional with scaled down exchange.
finegrid: boolean
Use fine grid for energy functional evaluations?
"""
if name == 'EXX':
assert hybrid is None and xc is None
hybrid = 1.0
xc = XC(XCNull())
elif name == 'PBE0':
assert hybrid is None and xc is None
hybrid = 0.25
xc = XC('HYB_GGA_XC_PBEH')
elif name == 'B3LYP':
assert hybrid is None and xc is None
hybrid = 0.2
xc = XC('HYB_GGA_XC_B3LYP')
if isinstance(xc, str):
xc = XC(xc)
self.hybrid = hybrid
self.xc = xc
self.type = xc.type
self.alpha = alpha
self.exx = 0.0
XCFunctional.__init__(self, name)
def get_setup_name(self):
return 'PBE'
def calculate_radial(self, rgd, n_sLg, Y_L, v_sg,
dndr_sLg=None, rnablaY_Lv=None,
tau_sg=None, dedtau_sg=None):
return self.xc.calculate_radial(rgd, n_sLg, Y_L, v_sg,
dndr_sLg, rnablaY_Lv)
def initialize(self, density, hamiltonian, wfs, occupations):
self.xc.initialize(density, hamiltonian, wfs, occupations)
self.nspins = wfs.nspins
self.setups = wfs.setups
self.density = density
self.kpt_u = wfs.kpt_u
self.gd = density.gd
self.kd = wfs.kd
self.bd = wfs.bd
N_c = self.gd.N_c
N = self.gd.N_c.prod()
vol = self.gd.dv * N
if self.alpha is None:
self.alpha = 6 * vol**(2 / 3.0) / pi**2
self.gamma = (vol / (2 * pi)**2 * sqrt(pi / self.alpha) *
self.kd.nbzkpts)
ecut = 0.5 * pi**2 / (self.gd.h_cv**2).sum(1).max()
if self.kd.N_c is None:
self.bzk_kc = np.zeros((1, 3))
dfghdfgh
else:
n = self.kd.N_c * 2 - 1
bzk_kc = np.indices(n).transpose((1, 2, 3, 0))
bzk_kc.shape = (-1, 3)
bzk_kc -= self.kd.N_c - 1
self.bzk_kc = bzk_kc.astype(float) / self.kd.N_c
self.pwd = PWDescriptor(ecut, self.gd, self.bzk_kc)
n = 0
for k_c, Gpk2_G in zip(self.bzk_kc[:], self.pwd.G2_qG):
if (k_c > -0.5).all() and (k_c <= 0.5).all(): #XXX???
if k_c.any():
self.gamma -= np.dot(np.exp(-self.alpha * Gpk2_G),
Gpk2_G**-1)
else:
self.gamma -= np.dot(np.exp(-self.alpha * Gpk2_G[1:]),
Gpk2_G[1:]**-1)
n += 1
assert n == self.kd.N_c.prod()
self.ghat = LFC(self.gd,
[setup.ghat_l for setup in density.setups],
dtype=complex
)
self.ghat.set_k_points(self.bzk_kc)
self.fullkd = KPointDescriptor(self.kd.bzk_kc, nspins=1)
class S:
id_a = []
def set_symmetry(self, s): pass
self.fullkd.set_symmetry(Atoms(pbc=True), S(), False)
self.fullkd.set_communicator(world)
self.pt = LFC(self.gd, [setup.pt_j for setup in density.setups],
dtype=complex)
self.pt.set_k_points(self.fullkd.ibzk_kc)
self.interpolator = density.interpolator
def set_positions(self, spos_ac):
self.ghat.set_positions(spos_ac)
self.pt.set_positions(spos_ac)
def calculate(self, gd, n_sg, v_sg=None, e_g=None):
# Normal XC contribution:
exc = self.xc.calculate(gd, n_sg, v_sg, e_g)
# Add EXX contribution:
return exc + self.exx
def calculate_exx(self):
"""Non-selfconsistent calculation."""
kd = self.kd
K = self.fullkd.nibzkpts
assert self.nspins == 1
Q = K // world.size
assert Q * world.size == K
parallel = (world.size > self.nspins)
self.exx = 0.0
self.exx_skn = np.zeros((self.nspins, K, self.bd.nbands))
kpt_u = []
for k in range(world.rank * Q, (world.rank + 1) * Q):
k_c = self.fullkd.ibzk_kc[k]
for k1, k1_c in enumerate(kd.bzk_kc):
if abs(k1_c - k_c).max() < 1e-10:
break
# Index of symmetry related point in the irreducible BZ
ik = kd.kibz_k[k1]
kpt = self.kpt_u[ik]
# KPoint from ground-state calculation
phase_cd = np.exp(2j * pi * self.gd.sdisp_cd * k_c[:, np.newaxis])
kpt2 = KPoint0(kpt.weight, kpt.s, k, None, phase_cd)
kpt2.psit_nG = np.empty_like(kpt.psit_nG)
kpt2.f_n = kpt.f_n / kpt.weight / K * 2
for n, psit_G in enumerate(kpt2.psit_nG):
psit_G[:] = kd.transform_wave_function(kpt.psit_nG[n], k1)
kpt2.P_ani = self.pt.dict(len(kpt.psit_nG))
self.pt.integrate(kpt2.psit_nG, kpt2.P_ani, k)
kpt_u.append(kpt2)
for s in range(self.nspins):
kpt1_q = [KPoint(self.fullkd, kpt) for kpt in kpt_u if kpt.s == s]
kpt2_q = kpt1_q[:]
if len(kpt1_q) == 0:
# No s-spins on this CPU:
continue
# Send rank:
srank = self.fullkd.get_rank_and_index(s, (kpt1_q[0].k - 1) % K)[0]
# Receive rank:
rrank = self.fullkd.get_rank_and_index(s, (kpt1_q[-1].k + 1) % K)[0]
# Shift k-points K // 2 times:
for i in range(K // 2 + 1):
if i < K // 2:
if parallel:
kpt = kpt2_q[-1].next()
kpt.start_receiving(rrank)
kpt2_q[0].start_sending(srank)
else:
kpt = kpt2_q[0]
for kpt1, kpt2 in zip(kpt1_q, kpt2_q):
if 2 * i == K:
self.apply(kpt1, kpt2, invert=(kpt1.k > kpt2.k))
else:
self.apply(kpt1, kpt2)
self.apply(kpt1, kpt2, invert=True)
if i < K // 2:
if parallel:
kpt.wait()
kpt2_q[0].wait()
kpt2_q.pop(0)
kpt2_q.append(kpt)
self.exx = world.sum(self.exx)
world.sum(self.exx_skn)
self.exx += self.calculate_paw_correction()
def apply(self, kpt1, kpt2, invert=False):
#print world.rank,kpt1.k,kpt2.k,invert
k1_c = self.fullkd.ibzk_kc[kpt1.k]
k2_c = self.fullkd.ibzk_kc[kpt2.k]
if invert:
k2_c = -k2_c
k12_c = k1_c - k2_c
N_c = self.gd.N_c
eikr_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k12_c / N_c).T)
for q, k_c in enumerate(self.bzk_kc):
if abs(k_c + k12_c).max() < 1e-9:
q0 = q
break
for q, k_c in enumerate(self.bzk_kc):
if abs(k_c - k12_c).max() < 1e-9:
q00 = q
break
Gpk2_G = self.pwd.G2_qG[q0]
if Gpk2_G[0] == 0:
Gpk2_G = Gpk2_G.copy()
Gpk2_G[0] = 1.0 / self.gamma
N = N_c.prod()
vol = self.gd.dv * N
nspins = self.nspins
same = (kpt1.k == kpt2.k)
for n1, psit1_R in enumerate(kpt1.psit_nG):
f1 = kpt1.f_n[n1]
for n2, psit2_R in enumerate(kpt2.psit_nG):
if same and n2 > n1:
continue
f2 = kpt2.f_n[n2]
nt_R = self.calculate_pair_density(n1, n2, kpt1, kpt2, q0,
invert)
nt_G = self.pwd.fft(nt_R * eikr_R) / N
vt_G = nt_G.copy()
vt_G *= -pi * vol / Gpk2_G
e = np.vdot(nt_G, vt_G).real * nspins * self.hybrid
if same and n1 == n2:
e /= 2
self.exx += e * f1 * f2
self.ekin -= 2 * e * f1 * f2
self.exx_skn[kpt1.s, kpt1.k, n1] += f2 * e
self.exx_skn[kpt2.s, kpt2.k, n2] += f1 * e
calculate_potential = not True
if calculate_potential:
vt_R = self.pwd.ifft(vt_G).conj() * eikr_R * N / vol
if kpt1 is kpt2 and not invert and n1 == n2:
kpt1.vt_nG[n1] = 0.5 * f1 * vt_R
if invert:
kpt1.Htpsit_nG[n1] += \
f2 * nspins * psit2_R.conj() * vt_R
else:
kpt1.Htpsit_nG[n1] += f2 * nspins * psit2_R * vt_R
if kpt1 is not kpt2:
if invert:
kpt2.Htpsit_nG[n2] += (f1 * nspins *
psit1_R.conj() * vt_R)
else:
kpt2.Htpsit_nG[n2] += (f1 * nspins *
psit1_R * vt_R.conj())
def calculate_paw_correction(self):
exx = 0
deg = 2 // self.nspins # spin degeneracy
for a, D_sp in self.density.D_asp.items():
setup = self.setups[a]
for D_p in D_sp:
D_ii = unpack2(D_p)
ni = len(D_ii)
for i1 in range(ni):
for i2 in range(ni):
A = 0.0
for i3 in range(ni):
p13 = packed_index(i1, i3, ni)
for i4 in range(ni):
p24 = packed_index(i2, i4, ni)
A += setup.M_pp[p13, p24] * D_ii[i3, i4]
p12 = packed_index(i1, i2, ni)
exx -= self.hybrid / deg * D_ii[i1, i2] * A
if setup.X_p is not None:
exx -= self.hybrid * np.dot(D_p, setup.X_p)
exx += self.hybrid * setup.ExxC
return exx
def calculate_pair_density(self, n1, n2, kpt1, kpt2, q, invert):
if invert:
nt_G = kpt1.psit_nG[n1].conj() * kpt2.psit_nG[n2].conj()
else:
nt_G = kpt1.psit_nG[n1].conj() * kpt2.psit_nG[n2]
Q_aL = {}
for a, P1_ni in kpt1.P_ani.items():
P1_i = P1_ni[n1]
P2_i = kpt2.P_ani[a][n2]
if invert:
D_ii = np.outer(P1_i.conj(), P2_i.conj())
else:
D_ii = np.outer(P1_i.conj(), P2_i)
D_p = pack(D_ii)
Q_aL[a] = np.dot(D_p, self.setups[a].Delta_pL)
self.ghat.add(nt_G, Q_aL, q)
return nt_G
class PhononCalculator:
"""This class defines the interface for phonon calculations."""
def __init__(self, calc, gamma=True, symmetry=False, e_ph=False,
communicator=serial_comm):
"""Inititialize class with a list of atoms.
The atoms object must contain a converged ground-state calculation.
The set of q-vectors in which the dynamical matrix will be calculated
is determined from the ``symmetry`` kwarg. For now, only time-reversal
symmetry is used to generate the irrecducible BZ.
Add a little note on parallelization strategy here.
Parameters
----------
calc: str or Calculator
Calculator containing a ground-state calculation.
gamma: bool
Gamma-point calculation with respect to the q-vector of the
dynamical matrix. When ``False``, the Monkhorst-Pack grid from the
ground-state calculation is used.
symmetry: bool
Use symmetries to reduce the q-vectors of the dynamcial matrix
(None, False or True). The different options are equivalent to the
options in a ground-state calculation.
e_ph: bool
Save the derivative of the effective potential.
communicator: Communicator
Communicator for parallelization over k-points and real-space
domain.
"""
# XXX
assert symmetry in [None, False], "Spatial symmetries not allowed yet"
self.symmetry = symmetry
if isinstance(calc, str):
self.calc = GPAW(calc, communicator=serial_comm, txt=None)
else:
self.calc = calc
# Make sure localized functions are initialized
self.calc.set_positions()
# Note that this under some circumstances (e.g. when called twice)
# allocates a new array for the P_ani coefficients !!
# Store useful objects
self.atoms = self.calc.get_atoms()
# Get rid of ``calc`` attribute
self.atoms.calc = None
# Boundary conditions
pbc_c = self.calc.atoms.get_pbc()
if np.all(pbc_c == False):
self.gamma = True
self.dtype = float
kpts = None
# Multigrid Poisson solver
poisson_solver = PoissonSolver()
else:
if gamma:
self.gamma = True
self.dtype = float
kpts = None
else:
self.gamma = False
self.dtype = complex
# Get k-points from ground-state calculation
kpts = self.calc.input_parameters.kpts
# FFT Poisson solver
poisson_solver = FFTPoissonSolver(dtype=self.dtype)
# K-point descriptor for the q-vectors of the dynamical matrix
# Note, no explicit parallelization here.
self.kd = KPointDescriptor(kpts, 1)
self.kd.set_symmetry(self.atoms, self.calc.wfs.setups,
usesymm=symmetry)
self.kd.set_communicator(serial_comm)
# Number of occupied bands
nvalence = self.calc.wfs.nvalence
nbands = nvalence / 2 + nvalence % 2
assert nbands <= self.calc.wfs.bd.nbands
# Extract other useful objects
# Ground-state k-point descriptor - used for the k-points in the
# ResponseCalculator
# XXX replace communicators when ready to parallelize
kd_gs = self.calc.wfs.kd
gd = self.calc.density.gd
kpt_u = self.calc.wfs.kpt_u
setups = self.calc.wfs.setups
dtype_gs = self.calc.wfs.dtype
# WaveFunctions
wfs = WaveFunctions(nbands, kpt_u, setups, kd_gs, gd, dtype=dtype_gs)
# Linear response calculator
self.response_calc = ResponseCalculator(self.calc, wfs,
dtype=self.dtype)
# Phonon perturbation
self.perturbation = PhononPerturbation(self.calc, self.kd,
poisson_solver,
dtype=self.dtype)
# Dynamical matrix
self.dyn = DynamicalMatrix(self.atoms, self.kd, dtype=self.dtype)
# Electron-phonon couplings
if e_ph:
self.e_ph = ElectronPhononCoupling(self.atoms, gd, self.kd,
dtype=self.dtype)
else:
self.e_ph = None
# Initialization flag
self.initialized = False
# Parallel communicator for parallelization over kpts and domain
self.comm = communicator
def initialize(self):
"""Initialize response calculator and perturbation."""
# Get scaled atomic positions
spos_ac = self.atoms.get_scaled_positions()
self.perturbation.initialize(spos_ac)
self.response_calc.initialize(spos_ac)
self.initialized = True
def __getstate__(self):
"""Method used when pickling.
Bound method attributes cannot be pickled and must therefore be deleted
before an instance is dumped to file.
"""
# Get state of object and take care of troublesome attributes
state = dict(self.__dict__)
state['kd'].__dict__['comm'] = serial_comm
state.pop('calc')
state.pop('perturbation')
state.pop('response_calc')
return state
def run(self, qpts_q=None, clean=False, name=None, path=None):
"""Run calculation for atomic displacements and update matrix.
Parameters
----------
qpts: List
List of q-points indices for which the dynamical matrix will be
calculated (only temporary).
"""
if not self.initialized:
self.initialize()
if self.gamma:
qpts_q = [0]
elif qpts_q is None:
qpts_q = range(self.kd.nibzkpts)
else:
assert isinstance(qpts_q, list)
# Update name and path attributes
self.set_name_and_path(name=name, path=path)
# Get string template for filenames
filename_str = self.get_filename_string()
# Delay the ranks belonging to the same k-point/domain decomposition
# equally
time.sleep(rank // self.comm.size)
# XXX Make a single ground_state_contributions member function
# Ground-state contributions to the force constants
self.dyn.density_ground_state(self.calc)
# self.dyn.wfs_ground_state(self.calc, self.response_calc)
# Calculate linear response wrt q-vectors and displacements of atoms
for q in qpts_q:
if not self.gamma:
self.perturbation.set_q(q)
# First-order contributions to the force constants
for a in self.dyn.indices:
for v in [0, 1, 2]:
# Check if the calculation has already been done
filename = filename_str % (q, a, v)
# Wait for all sub-ranks to enter
self.comm.barrier()
if os.path.isfile(os.path.join(self.path, filename)):
continue
if self.comm.rank == 0:
fd = open(os.path.join(self.path, filename), 'w')
# Wait for all sub-ranks here
self.comm.barrier()
components = ['x', 'y', 'z']
symbols = self.atoms.get_chemical_symbols()
print "q-vector index: %i" % q
print "Atom index: %i" % a
print "Atomic symbol: %s" % symbols[a]
print "Component: %s" % components[v]
# Set atom and cartesian component of perturbation
self.perturbation.set_av(a, v)
# Calculate linear response
self.response_calc(self.perturbation)
# Calculate row of the matrix of force constants
self.dyn.calculate_row(self.perturbation,
self.response_calc)
# Write force constants to file
if self.comm.rank == 0:
self.dyn.write(fd, q, a, v)
fd.close()
# Store effective potential derivative
if self.e_ph is not None:
v1_eff_G = self.perturbation.v1_G + \
self.response_calc.vHXC1_G
self.e_ph.v1_eff_qavG.append(v1_eff_G)
# Wait for the file-writing rank here
self.comm.barrier()
# XXX
# Check that all files are valid and collect in a single file
# Remove the files
if clean:
self.clean()
def get_atoms(self):
"""Return atoms."""
return self.atoms
def get_dynamical_matrix(self):
"""Return reference to ``dyn`` attribute."""
return self.dyn
def get_filename_string(self):
"""Return string template for force constant filenames."""
name_str = (self.name + '.' + 'q_%%0%ii_' % len(str(self.kd.nibzkpts)) +
'a_%%0%ii_' % len(str(len(self.atoms))) + 'v_%i' + '.pckl')
return name_str
def set_atoms(self, atoms):
"""Set atoms to be included in the calculation.
Parameters
----------
atoms: list
Can be either a list of strings, ints or ...
"""
assert isinstance(atoms, list)
if isinstance(atoms[0], str):
assert np.all([isinstance(atom, str) for atom in atoms])
sym_a = self.atoms.get_chemical_symbols()
# List for atomic indices
indices = []
for type in atoms:
indices.extend([a for a, atom in enumerate(sym_a)
if atom == type])
else:
assert np.all([isinstance(atom, int) for atom in atoms])
indices = atoms
self.dyn.set_indices(indices)
def set_name_and_path(self, name=None, path=None):
"""Set name and path of the force constant files.
name: str
Base name for the files which the elements of the matrix of force
constants will be written to.
path: str
Path specifying the directory where the files will be dumped.
"""
if name is None:
self.name = 'phonon.' + self.atoms.get_chemical_formula()
else:
self.name = name
# self.name += '.nibzkpts_%i' % self.kd.nibzkpts
if path is None:
self.path = '.'
else:
self.path = path
# Set corresponding attributes in the ``dyn`` attribute
filename_str = self.get_filename_string()
self.dyn.set_name_and_path(filename_str, self.path)
def clean(self):
"""Delete generated files."""
filename_str = self.get_filename_string()
for q in range(self.kd.nibzkpts):
for a in range(len(self.atoms)):
for v in [0, 1, 2]:
filename = filename_str % (q, a, v)
if os.path.isfile(os.path.join(self.path, filename)):
os.remove(filename)
def band_structure(self, path_kc, modes=False, acoustic=True):
"""Calculate phonon dispersion along a path in the Brillouin zone.
The dynamical matrix at arbitrary q-vectors is obtained by Fourier
transforming the real-space matrix. In case of negative eigenvalues
(squared frequency), the corresponding negative frequency is returned.
Parameters
----------
path_kc: ndarray
List of k-point coordinates (in units of the reciprocal lattice
vectors) specifying the path in the Brillouin zone for which the
dynamical matrix will be calculated.
modes: bool
Returns both frequencies and modes (mass scaled) when True.
acoustic: bool
Restore the acoustic sum-rule in the calculated force constants.
"""
for k_c in path_kc:
assert np.all(np.asarray(k_c) <= 1.0), \
"Scaled coordinates must be given"
# Assemble the dynanical matrix from calculated force constants
self.dyn.assemble(acoustic=acoustic)
# Get the dynamical matrix in real-space
DR_lmn, R_clmn = self.dyn.real_space()
# Reshape for the evaluation of the fourier sums
shape = DR_lmn.shape
DR_m = DR_lmn.reshape((-1,) + shape[-2:])
R_cm = R_clmn.reshape((3, -1))
# Lists for frequencies and modes along path
omega_kn = []
u_kn = []
# Number of atoms included
N = len(self.dyn.get_indices())
# Mass prefactor for the normal modes
m_inv_av = self.dyn.get_mass_array()
for q_c in path_kc:
# Evaluate fourier transform
phase_m = np.exp(-2.j * pi * np.dot(q_c, R_cm))
# Dynamical matrix in unit of Ha / Bohr**2 / amu
D_q = np.sum(phase_m[:, np.newaxis, np.newaxis] * DR_m, axis=0)
if modes:
omega2_n, u_avn = la.eigh(D_q, UPLO='L')
# Sort eigenmodes according to eigenvalues (see below) and
# multiply with mass prefactor
u_nav = u_avn[:, omega2_n.argsort()].T.copy() * m_inv_av
# Multiply with mass prefactor
u_kn.append(u_nav.reshape((3*N, -1, 3)))
else:
omega2_n = la.eigvalsh(D_q, UPLO='L')
# Sort eigenvalues in increasing order
omega2_n.sort()
# Use dtype=complex to handle negative eigenvalues
omega_n = np.sqrt(omega2_n.astype(complex))
# Take care of imaginary frequencies
if not np.all(omega2_n >= 0.):
indices = np.where(omega2_n < 0)[0]
print ("WARNING, %i imaginary frequencies at "
"q = (% 5.2f, % 5.2f, % 5.2f) ; (omega_q =% 5.3e*i)"
% (len(indices), q_c[0], q_c[1], q_c[2],
omega_n[indices][0].imag))
omega_n[indices] = -1 * np.sqrt(np.abs(omega2_n[indices].real))
omega_kn.append(omega_n.real)
# Conversion factor from sqrt(Ha / Bohr**2 / amu) -> eV
s = units.Hartree**0.5 * units._hbar * 1.e10 / \
(units._e * units._amu)**(0.5) / units.Bohr
# Convert to eV and Ang
omega_kn = s * np.asarray(omega_kn)
if modes:
u_kn = np.asarray(u_kn) * units.Bohr
return omega_kn, u_kn
return omega_kn
def write_modes(self, q_c, branches=0, kT=units.kB*300, repeat=(1, 1, 1),
nimages=30, acoustic=True):
"""Write mode to trajectory file.
The classical equipartioning theorem states that each normal mode has
an average energy::
<E> = 1/2 * k_B * T = 1/2 * omega^2 * Q^2
=>
Q = sqrt(k_B*T) / omega
at temperature T. Here, Q denotes the normal coordinate of the mode.
Parameters
----------
q_c: ndarray
q-vector of the modes.
branches: int or list
Branch index of calculated modes.
kT: float
Temperature in units of eV. Determines the amplitude of the atomic
displacements in the modes.
repeat: tuple
Repeat atoms (l, m, n) times in the directions of the lattice
vectors. Displacements of atoms in repeated cells carry a Bloch
phase factor given by the q-vector and the cell lattice vector R_m.
nimages: int
Number of images in an oscillation.
"""
if isinstance(branches, int):
branch_n = [branches]
else:
branch_n = list(branches)
# Calculate modes
omega_n, u_n = self.band_structure([q_c], modes=True, acoustic=acoustic)
# Repeat atoms
atoms = self.atoms * repeat
pos_mav = atoms.positions.copy()
# Total number of unit cells
M = np.prod(repeat)
# Corresponding lattice vectors R_m
R_cm = np.indices(repeat[::-1]).reshape(3, -1)[::-1]
# Bloch phase
phase_m = np.exp(2.j * pi * np.dot(q_c, R_cm))
phase_ma = phase_m.repeat(len(self.atoms))
for n in branch_n:
omega = omega_n[0, n]
u_av = u_n[0, n] # .reshape((-1, 3))
# Mean displacement at high T ?
u_av *= sqrt(kT / abs(omega))
mode_av = np.zeros((len(self.atoms), 3), dtype=self.dtype)
indices = self.dyn.get_indices()
mode_av[indices] = u_av
mode_mav = (np.vstack([mode_av]*M) * phase_ma[:, np.newaxis]).real
traj = PickleTrajectory('%s.mode.%d.traj' % (self.name, n), 'w')
for x in np.linspace(0, 2*pi, nimages, endpoint=False):
# XXX Is it correct to take out the sine component here ?
atoms.set_positions(pos_mav + sin(x) * mode_mav)
traj.write(atoms)
traj.close()
class UTDomainParallelSetup(TestCase):
"""
Setup a simple domain parallel calculation."""
# Number of bands
nbands = 12
# Spin-polarized
nspins = 1
# Mean spacing and number of grid points per axis (G x G x G)
h = 0.25 / Bohr
G = 48
# Type of boundary conditions employed (determines nibzkpts and dtype)
boundaries = None
nibzkpts = None
dtype = None
timer = nulltimer
def setUp(self):
for virtvar in ['boundaries']:
assert getattr(self,
virtvar) is not None, 'Virtual "%s"!' % virtvar
# Basic unit cell information:
res, N_c = shapeopt(100, self.G**3, 3, 0.2)
#N_c = 4*np.round(np.array(N_c)/4) # makes domain decomposition easier
cell_cv = self.h * np.diag(N_c)
pbc_c = {'zero' : (False,False,False), \
'periodic': (True,True,True), \
'mixed' : (True, False, True)}[self.boundaries]
# Create randomized gas-like atomic configuration on interim grid
tmpgd = GridDescriptor(N_c, cell_cv, pbc_c)
self.atoms = create_random_atoms(tmpgd)
# Create setups
Z_a = self.atoms.get_atomic_numbers()
assert 1 == self.nspins
self.setups = Setups(Z_a, p.setups, p.basis, p.lmax, xc)
self.natoms = len(self.setups)
# Decide how many kpoints to sample from the 1st Brillouin Zone
kpts_c = np.ceil(
(10 / Bohr) / np.sum(cell_cv**2, axis=1)**0.5).astype(int)
kpts_c = tuple(kpts_c * pbc_c + 1 - pbc_c)
self.bzk_kc = kpts2ndarray(kpts_c)
# Set up k-point descriptor
self.kd = KPointDescriptor(self.bzk_kc, self.nspins)
self.kd.set_symmetry(self.atoms, self.setups, p.usesymm)
# Set the dtype
if self.kd.gamma:
self.dtype = float
else:
self.dtype = complex
# Create communicators
parsize, parsize_bands = self.get_parsizes()
assert self.nbands % np.prod(parsize_bands) == 0
domain_comm, kpt_comm, band_comm = distribute_cpus(
parsize, parsize_bands, self.nspins, self.kd.nibzkpts)
self.kd.set_communicator(kpt_comm)
# Set up band descriptor:
self.bd = BandDescriptor(self.nbands, band_comm)
# Set up grid descriptor:
self.gd = GridDescriptor(N_c, cell_cv, pbc_c, domain_comm, parsize)
# Set up kpoint/spin descriptor (to be removed):
self.kd_old = KPointDescriptorOld(self.nspins, self.kd.nibzkpts,
kpt_comm, self.kd.gamma, self.dtype)
def tearDown(self):
del self.atoms, self.bd, self.gd, self.kd, self.kd_old
def get_parsizes(self):
# Careful, overwriting imported GPAW params may cause amnesia in Python.
from gpaw import parsize, parsize_bands
# D: number of domains
# B: number of band groups
if parsize is None:
D = min(world.size, 2)
else:
D = parsize
assert world.size % D == 0
if parsize_bands is None:
B = world.size // D
else:
B = parsize_bands
return D, B
# =================================
def verify_comm_sizes(self):
if world.size == 1:
return
comm_sizes = tuple([comm.size for comm in [world, self.bd.comm, \
self.gd.comm, self.kd_old.comm]])
self._parinfo = '%d world, %d band, %d domain, %d kpt' % comm_sizes
self.assertEqual(self.nbands % self.bd.comm.size, 0)
self.assertEqual(
(self.nspins * self.kd.nibzkpts) % self.kd_old.comm.size, 0)
def ibz2bz(self, atoms):
"""Transform wave functions in IBZ to the full BZ."""
assert self.kd.comm.size == 1
# New k-point descriptor for full BZ:
kd = KPointDescriptor(self.kd.bzk_kc, nspins=self.nspins)
kd.set_communicator(serial_comm)
self.pt = LFC(self.gd, [setup.pt_j for setup in self.setups],
kd,
dtype=self.dtype)
self.pt.set_positions(atoms.get_scaled_positions())
self.initialize_wave_functions_from_restart_file()
weight = 2.0 / kd.nspins / kd.nbzkpts
# Build new list of k-points:
kpt_u = []
for s in range(self.nspins):
for k in range(kd.nbzkpts):
# Index of symmetry related point in the IBZ
ik = self.kd.bz2ibz_k[k]
r, u = self.kd.get_rank_and_index(s, ik)
assert r == 0
kpt = self.mykpts[u]
phase_cd = np.exp(2j * np.pi * self.gd.sdisp_cd *
kd.bzk_kc[k, :, np.newaxis])
# New k-point:
kpt2 = KPoint(weight, s, k, k, phase_cd)
kpt2.f_n = kpt.f_n / kpt.weight / kd.nbzkpts * 2 / self.nspins
kpt2.eps_n = kpt.eps_n.copy()
# Transform wave functions using symmetry operation:
Psit_nG = self.gd.collect(kpt.psit_nG)
if Psit_nG is not None:
Psit_nG = Psit_nG.copy()
for Psit_G in Psit_nG:
Psit_G[:] = self.kd.transform_wave_function(Psit_G, k)
kpt2.psit = UniformGridWaveFunctions(self.bd.nbands,
self.gd,
self.dtype,
kpt=k,
dist=(self.bd.comm,
self.bd.comm.size),
spin=kpt.s,
collinear=True)
self.gd.distribute(Psit_nG, kpt2.psit_nG)
# Calculate PAW projections:
nproj_a = [setup.ni for setup in self.setups]
kpt2.projections = Projections(self.bd.nbands,
nproj_a,
kpt.projections.atom_partition,
self.bd.comm,
collinear=True,
spin=s,
dtype=self.dtype)
kpt2.psit.matrix_elements(self.pt, out=kpt2.projections)
kpt_u.append(kpt2)
self.kd = kd
self.mykpts = kpt_u
