def test_something(self):
laplace_uG = np.empty_like(self.laplace0_uG)
op = Laplace(self.gd, dtype=self.dtype)
for myu, laplace_G in enumerate(laplace_uG):
phase_cd = {float:None, complex:self.phase_ucd[myu]}[self.dtype]
op.apply(self.wf_uG[myu], laplace_G, phase_cd)
print 'myu:', myu, 'diff:', np.std(laplace_G-self.laplace0_uG[myu]), '/', np.abs(laplace_G-self.laplace0_uG[myu]).max()
def test_something(self):
laplace_uG = np.empty_like(self.laplace0_uG)
op = Laplace(self.gd, dtype=self.dtype)
for myu, laplace_G in enumerate(laplace_uG):
phase_cd = {float:None, complex:self.phase_ucd[myu]}[self.dtype]
op.apply(self.wf_uG[myu], laplace_G, phase_cd)
print 'myu:', myu, 'diff:', np.std(laplace_G-self.laplace0_uG[myu]), '/', np.abs(laplace_G-self.laplace0_uG[myu]).max()
def go(comm, ngpts, repeat, narrays, out, prec):
N_c = np.array((ngpts, ngpts, ngpts))
a = 10.0
gd = GridDescriptor(N_c, (a, a, a), comm=comm))
gdcoarse = gd.coarsen()
gdfine = gd.refine()
kin1 = Laplace(gd, -0.5, 1).apply
laplace = Laplace(gd, -0.5, 2)
kin2 = laplace.apply
restrict = Transformer(gd, gdcoarse, 1).apply
interpolate = Transformer(gd, gdfine, 1).apply
precondition = Preconditioner(gd, laplace, np_float)
a1 = gd.empty(narrays)
a1[:] = 1.0
a2 = gd.empty(narrays)
c = gdcoarse.empty(narrays)
f = gdfine.empty(narrays)
T = [0, 0, 0, 0, 0]
for i in range(repeat):
comm.barrier()
kin1(a1, a2)
comm.barrier()
t0a = time()
kin1(a1, a2)
t0b = time()
comm.barrier()
t1a = time()
kin2(a1, a2)
t1b = time()
comm.barrier()
t2a = time()
for A, C in zip(a1,c):
restrict(A, C)
t2b = time()
comm.barrier()
t3a = time()
for A, F in zip(a1,f):
interpolate(A, F)
t3b = time()
comm.barrier()
if prec:
t4a = time()
for A in a1:
precondition(A, None, None, None)
t4b = time()
comm.barrier()
T[0] += t0b - t0a
T[1] += t1b - t1a
T[2] += t2b - t2a
T[3] += t3b - t3a
if prec:
T[4] += t4b - t4a
if mpi.rank == 0:
out.write(' %2d %2d %2d' % tuple(gd.parsize_c))
out.write(' %12.6f %12.6f %12.6f %12.6f %12.6f\n' %
tuple([t / repeat / narrays for t in T]))
out.flush()
def __init__(self, hamiltonian, wfs, gd, dtype=float):
"""Save useful objects for the Sternheimer operator.
Parameters
----------
hamiltonian: Hamiltonian
Hamiltonian for a ground-state calculation.
wfs: WaveFunctions
Ground-state wave-functions.
gd: GridDescriptor
Grid on which the operator is defined.
dtype: dtype
dtype of the wave-function being operated on.
"""
self.hamiltonian = hamiltonian
self.kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
self.kpt_u = wfs.kpt_u
self.pt = wfs.pt
self.gd = gd
# Variables for k-point and band index
self.k = None
self.n = None
self.kplusq = None
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N, N)
self.dtype = dtype
def __init__(self,
stencil,
parallel,
initksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
kptband_comm,
timer,
reuse_wfs_method=None,
collinear=True):
FDPWWaveFunctions.__init__(self,
parallel,
initksl,
reuse_wfs_method=reuse_wfs_method,
collinear=collinear,
gd=gd,
nvalence=nvalence,
setups=setups,
bd=bd,
dtype=dtype,
world=world,
kd=kd,
kptband_comm=kptband_comm,
timer=timer)
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype)
self.taugrad_v = None # initialized by MGGA functional
def set_grid_descriptor(self, gd):
# Should probably be renamed initialize
self.gd = gd
scale = -0.25 / pi
if self.nn == 'M':
if not gd.orthogonal:
raise RuntimeError('Cannot use Mehrstellen stencil with '
'non orthogonal cell.')
self.operators = [LaplaceA(gd, -scale)]
self.B = LaplaceB(gd)
else:
self.operators = [Laplace(gd, scale, self.nn)]
self.B = None
self.interpolators = []
self.restrictors = []
level = 0
self.presmooths = [2]
self.postsmooths = [1]
# Weights for the relaxation,
# only used if 'J' (Jacobi) is chosen as method
self.weights = [2.0 / 3.0]
while level < 8:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1))
self.interpolators.append(Transformer(gd2, gd))
self.restrictors.append(Transformer(gd, gd2))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
if self.operators[-1].gd.N_c.max() > 36:
# Try to warn exactly once no matter how one uses the solver.
if gd.comm.parent is None:
warn = (gd.comm.rank == 0)
else:
warn = (gd.comm.parent.rank == 0)
if warn:
warntxt = '\n'.join([POISSON_GRID_WARNING, '',
self.get_description()])
else:
warntxt = ('Poisson warning from domain rank %d'
% self.gd.comm.rank)
# Warn from all ranks to avoid deadlocks.
warnings.warn(warntxt, stacklevel=2)
def apply_t(self):
"""Apply kinetic energy operator and return new object."""
p = 2 # padding
newsize_c = self.size_c + 2 * p
gd = GridDescriptor(N_c=newsize_c + 1, cell_cv=self.gd.h_c * (newsize_c + 1), pbc_c=False, comm=mpi.serial_comm)
T = Laplace(gd, scale=1 / 2.0, n=p)
f_ig = np.zeros((len(self.f_iG),) + tuple(newsize_c))
f_ig[:, p:-p, p:-p, p:-p] = self.f_iG
Tf_iG = np.empty_like(f_ig)
T.apply(f_ig, Tf_iG)
return LocalizedFunctions(self.gd, Tf_iG, self.corner_c - p, self.index)
def __init__(self, stencil, diagksl, orthoksl, initksl, gd, nvalence,
setups, bd, dtype, world, kd, kptband_comm, timer):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, dtype, world, kd,
kptband_comm, timer)
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype)
self.matrixoperator = MatrixOperator(self.orthoksl)
self.taugrad_v = None # initialized by MGGA functional
def apply_t(self):
"""Apply kinetic energy operator and return new object."""
p = 2 # padding
newsize_c = self.size_c + 2 * p
gd = GridDescriptor(N_c=newsize_c + 1,
cell_cv=self.gd.h_c * (newsize_c + 1),
pbc_c=False,
comm=mpi.serial_comm)
T = Laplace(gd, scale =1/2., n=p)
f_ig = np.zeros((len(self.f_iG),) + tuple(newsize_c))
f_ig[:, p:-p, p:-p, p:-p] = self.f_iG
Tf_iG = np.empty_like(f_ig)
T.apply(f_ig, Tf_iG)
return LocalizedFunctions(self.gd, Tf_iG, self.corner_c - p,
self.index)
def __init__(self, hamiltonian, wfs, gd, dtype=float):
"""Save useful objects for the Sternheimer operator.
Parameters
----------
hamiltonian: Hamiltonian
Hamiltonian for a ground-state calculation.
wfs: WaveFunctions
Ground-state wave-functions.
gd: GridDescriptor
Grid on which the operator is defined.
dtype: dtype
dtype of the wave-function being operated on.
"""
self.hamiltonian = hamiltonian
self.kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
self.kpt_u = wfs.kpt_u
self.pt = wfs.pt
self.gd = gd
# Variables for k-point and band index
self.k = None
self.n = None
self.kplusq = None
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N, N)
self.dtype = dtype
def set_grid_descriptor(self, gd):
# Should probably be renamed initialize
self.gd = gd
self.dv = gd.dv
gd = self.gd
scale = -0.25 / pi
if self.nn == 'M':
if not gd.orthogonal:
raise RuntimeError('Cannot use Mehrstellen stencil with '
'non orthogonal cell.')
self.operators = [LaplaceA(gd, -scale, allocate=False)]
self.B = LaplaceB(gd, allocate=False)
else:
self.operators = [Laplace(gd, scale, self.nn, allocate=False)]
self.B = None
self.interpolators = []
self.restrictors = []
level = 0
self.presmooths = [2]
self.postsmooths = [1]
# Weights for the relaxation,
# only used if 'J' (Jacobi) is chosen as method
self.weights = [2.0 / 3.0]
while level < 4:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1, allocate=False))
self.interpolators.append(Transformer(gd2, gd, allocate=False))
self.restrictors.append(Transformer(gd, gd2, allocate=False))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
def __init__(self, gd0, kin0, dtype=float, block=1):
gd1 = gd0.coarsen()
gd2 = gd1.coarsen()
self.kin0 = kin0
self.kin1 = Laplace(gd1, -0.5, 1, dtype)
self.kin2 = Laplace(gd2, -0.5, 1, dtype)
self.scratch0 = gd0.zeros((2, block), dtype, False)
self.scratch1 = gd1.zeros((3, block),dtype, False)
self.scratch2 = gd2.zeros((3, block), dtype, False)
self.step = 0.66666666 / kin0.get_diagonal_element()
self.restrictor_object0 = Transformer(gd0, gd1, 1,dtype)
self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype)
self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype)
self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype)
self.restrictor0 = self.restrictor_object0.apply
self.restrictor1 = self.restrictor_object1.apply
self.interpolator2 = self.interpolator_object2.apply
self.interpolator1 = self.interpolator_object1.apply
def __init__(self, gd0, kin0, dtype=float, block=1):
gd1 = gd0.coarsen()
gd2 = gd1.coarsen()
self.kin0 = kin0
self.kin1 = Laplace(gd1, -0.5, 1, dtype)
self.kin2 = Laplace(gd2, -0.5, 1, dtype)
self.scratch0 = gd0.zeros((2, block), dtype, False)
self.scratch1 = gd1.zeros((3, block), dtype, False)
self.scratch2 = gd2.zeros((3, block), dtype, False)
self.step = 0.66666666 / kin0.get_diagonal_element()
self.restrictor_object0 = Transformer(gd0, gd1, 1, dtype)
self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype)
self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype)
self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype)
self.restrictor0 = self.restrictor_object0.apply
self.restrictor1 = self.restrictor_object1.apply
self.interpolator2 = self.interpolator_object2.apply
self.interpolator1 = self.interpolator_object1.apply
def __init__(self, gd, bd, kd, setups, dtype): # override constructor
assert kd.comm.size == 1
WaveFunctions.__init__(self, gd, 1, setups, bd, dtype, world, kd, None)
self.kin = Laplace(gd, -0.5, dtype=dtype, allocate=False)
self.diagksl = None
self.orthoksl = BandLayouts(gd, bd, dtype)
self.initksl = None
self.overlap = None
self.rank_a = None
class Preconditioner:
def __init__(self, gd0, kin0, dtype=float, block=1):
gd1 = gd0.coarsen()
gd2 = gd1.coarsen()
self.kin0 = kin0
self.kin1 = Laplace(gd1, -0.5, 1, dtype)
self.kin2 = Laplace(gd2, -0.5, 1, dtype)
self.scratch0 = gd0.zeros((2, block), dtype, False)
self.scratch1 = gd1.zeros((3, block), dtype, False)
self.scratch2 = gd2.zeros((3, block), dtype, False)
self.step = 0.66666666 / kin0.get_diagonal_element()
self.restrictor_object0 = Transformer(gd0, gd1, 1, dtype)
self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype)
self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype)
self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype)
self.restrictor0 = self.restrictor_object0.apply
self.restrictor1 = self.restrictor_object1.apply
self.interpolator2 = self.interpolator_object2.apply
self.interpolator1 = self.interpolator_object1.apply
def calculate_kinetic_energy(self, psit_xG, kpt):
return None
def __call__(self, residuals, kpt, ekin=None, out=None):
if residuals.ndim == 3:
if out is None:
return self.__call__(residuals[np.newaxis], kpt)[0]
return self.__call__(residuals[np.newaxis], kpt,
out=out[np.newaxis])[0]
nb = len(residuals) # number of bands
phases = kpt.phase_cd
step = self.step
if out is None:
d0, q0 = self.scratch0[:, :nb]
else:
d0 = out
q0 = self.scratch0[0, :nb]
r1, d1, q1 = self.scratch1[:, :nb]
r2, d2, q2 = self.scratch2[:, :nb]
self.restrictor0(-residuals, r1, phases)
d1[:] = 4 * step * r1
self.kin1.apply(d1, q1, phases)
q1 -= r1
self.restrictor1(q1, r2, phases)
d2 = 16 * step * r2
self.kin2.apply(d2, q2, phases)
q2 -= r2
d2 -= 16 * step * q2
self.interpolator2(d2, q1, phases)
d1 -= q1
self.kin1.apply(d1, q1, phases)
q1 -= r1
d1 -= 4 * step * q1
self.interpolator1(-d1, d0, phases)
self.kin0.apply(d0, q0, phases)
q0 -= residuals
axpy(-step, q0, d0) # d0 -= step * q0
d0 *= -1.0
return d0
def __init__(self, stencil, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer)
self.wd = self.gd # wave function descriptor
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype, allocate=False)
self.matrixoperator = MatrixOperator(orthoksl)
def __init__(self, stencil, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer)
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype)
self.matrixoperator = MatrixOperator(self.orthoksl)
self.taugrad_v = None # initialized by MGGA functional
class Preconditioner:
def __init__(self, gd0, kin0, dtype=float, block=1):
gd1 = gd0.coarsen()
gd2 = gd1.coarsen()
self.kin0 = kin0
self.kin1 = Laplace(gd1, -0.5, 1, dtype)
self.kin2 = Laplace(gd2, -0.5, 1, dtype)
self.scratch0 = gd0.zeros((2, block), dtype, False)
self.scratch1 = gd1.zeros((3, block), dtype, False)
self.scratch2 = gd2.zeros((3, block), dtype, False)
self.step = 0.66666666 / kin0.get_diagonal_element()
self.restrictor_object0 = Transformer(gd0, gd1, 1, dtype, False)
self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype, False)
self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype, False)
self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype, False)
self.restrictor0 = self.restrictor_object0.apply
self.restrictor1 = self.restrictor_object1.apply
self.interpolator2 = self.interpolator_object2.apply
self.interpolator1 = self.interpolator_object1.apply
self.allocated = False
def allocate(self):
assert not self.allocated
for transformer in [
self.restrictor_object0, self.restrictor_object1,
self.interpolator_object2, self.interpolator_object1
]:
transformer.allocate()
self.allocated = True
def __call__(self, residuals, kpt):
nb = len(residuals) # number of bands
phases = kpt.phase_cd
step = self.step
d0, q0 = self.scratch0[:, :nb]
r1, d1, q1 = self.scratch1[:, :nb]
r2, d2, q2 = self.scratch2[:, :nb]
self.restrictor0(-residuals, r1, phases)
d1[:] = 4 * step * r1
self.kin1.apply(d1, q1, phases)
q1 -= r1
self.restrictor1(q1, r2, phases)
d2 = 16 * step * r2
self.kin2.apply(d2, q2, phases)
q2 -= r2
d2 -= 16 * step * q2
self.interpolator2(d2, q1, phases)
d1 -= q1
self.kin1.apply(d1, q1, phases)
q1 -= r1
d1 -= 4 * step * q1
self.interpolator1(-d1, d0, phases)
self.kin0.apply(d0, q0, phases)
q0 -= residuals
axpy(-step, q0, d0) # d0 -= step * q0
d0 *= -1.0
return d0
class Preconditioner:
def __init__(self, gd0, kin0, dtype=float, block=1):
gd1 = gd0.coarsen()
gd2 = gd1.coarsen()
self.kin0 = kin0
self.kin1 = Laplace(gd1, -0.5, 1, dtype)
self.kin2 = Laplace(gd2, -0.5, 1, dtype)
self.scratch0 = gd0.zeros((2, block), dtype, False)
self.scratch1 = gd1.zeros((3, block),dtype, False)
self.scratch2 = gd2.zeros((3, block), dtype, False)
self.step = 0.66666666 / kin0.get_diagonal_element()
self.restrictor_object0 = Transformer(gd0, gd1, 1,dtype, False)
self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype, False)
self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype, False)
self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype, False)
self.restrictor0 = self.restrictor_object0.apply
self.restrictor1 = self.restrictor_object1.apply
self.interpolator2 = self.interpolator_object2.apply
self.interpolator1 = self.interpolator_object1.apply
self.allocated = False
def allocate(self):
assert not self.allocated
for transformer in [self.restrictor_object0,
self.restrictor_object1,
self.interpolator_object2,
self.interpolator_object1]:
transformer.allocate()
self.allocated = True
def __call__(self, residuals, kpt):
nb = len(residuals) # number of bands
phases = kpt.phase_cd
step = self.step
d0, q0 = self.scratch0[:,:nb]
r1, d1, q1 = self.scratch1[:, :nb]
r2, d2, q2 = self.scratch2[:, :nb]
self.restrictor0(-residuals, r1, phases)
d1[:] = 4 * step * r1
self.kin1.apply(d1, q1, phases)
q1 -= r1
self.restrictor1(q1, r2, phases)
d2 = 16 * step * r2
self.kin2.apply(d2, q2, phases)
q2 -= r2
d2 -= 16 * step * q2
self.interpolator2(d2, q1, phases)
d1 -= q1
self.kin1.apply(d1, q1, phases)
q1 -= r1
d1 -= 4 * step * q1
self.interpolator1(-d1, d0, phases)
self.kin0.apply(d0, q0, phases)
q0 -= residuals
axpy(-step, q0, d0) # d0 -= step * q0
d0 *= -1.0
return d0
def __init__(self,
stencil,
diagksl,
orthoksl,
initksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, dtype, world, kd,
timer)
self.wd = self.gd # wave function descriptor
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype, allocate=False)
self.matrixoperator = MatrixOperator(orthoksl)
def set_absorbing_boundary(self, absorbing_boundary):
""" Sets up the absorbing boundary.
Parameters:
absorbing_boundary: absorbing boundary object of any kind.
"""
self.absorbing_boundary = absorbing_boundary
self.absorbing_boundary.set_up(self.hamiltonian.gd)
if self.absorbing_boundary.type == 'PML':
gd = self.hamiltonian.gd
self.laplace = Laplace(gd, n=2, dtype=complex)
self.gradient = np.array(
(Gradient(gd, 0, n=2,
dtype=complex), Gradient(gd, 1, n=2, dtype=complex),
Gradient(gd, 2, n=2, dtype=complex)))
self.lpsit = None
def initialize(self, load_gauss=False):
self.presmooths[self.levels] = 8
self.postsmooths[self.levels] = 8
self.phis = [None] + [gd.zeros() for gd in self.gds[1:]]
self.residuals = [gd.zeros() for gd in self.gds]
self.rhos = [gd.zeros() for gd in self.gds]
self.op_coarse_weights = [[g.empty() for g in (gd, ) * 4]
for gd in self.gds[1:]]
scale = -0.25 / np.pi
for i, gd in enumerate(self.gds):
if i == 0:
nn = self.nn
weights = self.dielectric.eps_gradeps
else:
nn = 1
weights = self.op_coarse_weights[i - 1]
operators = [Laplace(gd, scale, nn)] + \
[Gradient(gd, j, scale, nn) for j in (0, 1, 2)]
self.operators.append(WeightedFDOperator(weights, operators))
if load_gauss:
self.load_gauss()
def __init__(self, gd, project, dtype=float):
"""Init the gpaw preconditioner.
Parameters
----------
gd: GridDescriptor
Coarse grid
"""
self.project = project
self.gd = gd
# K-point for the preconditioner
self.kpt = None
kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
self.pc = Preconditioner(gd, kin, dtype=dtype)
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N,N)
self.dtype = dtype
def create_laplace(self, gd, scale=1.0, n=1, dtype=float):
"""Instantiate and return a Laplace operator
Allows subclasses to change the Laplace operator
"""
return Laplace(gd, scale, n, dtype)
def create_laplace(self, gd, scale=1.0, n=1, dtype=float):
operators = [Laplace(gd, scale, n, dtype)]
operators += [Gradient(gd, j, scale, n, dtype) for j in (0, 1, 2)]
return WeightedFDOperator(operators)
def initialize(self, dens, ham, wfs, occ):
MGGA.initialize(self, dens, ham, wfs, occ)
self.kernel.world = wfs.world
self.kernel.gd = dens.finegd
self.kernel.lapl = Laplace(dens.finegd)
from time import time
from gpaw.grid_descriptor import GridDescriptor
from gpaw.fd_operators import Laplace
import gpaw.mpi as mpi
n = 96
h = 0.1
L = n * h
gd = GridDescriptor((n, n, n), [L, L, L])
# Allocate arrays:
a = gd.zeros(100) + 1.2
b = gd.empty(100)
o = Laplace(gd, 2).apply
t0 = time()
for r in range(10):
o(a, b)
if mpi.rank == 0:
print time() - t0, a.shape
offload_enabled = os.environ.get("GPAW_OFFLOAD");
device = mic.devices[0]
gpts = 64
gd = GridDescriptor([gpts, gpts, gpts])
a = gd.empty() # source
b = gd.empty() # result
np.random.seed(10)
a[:] = np.random.random(a.shape)
b[:] = np.zeros(b.shape)
op = Laplace(gd, 1.0, 3)
if offload_enabled:
offl_a = device.bind(a)
offl_b = device.bind(b)
print "--------------------------------------------------------------"
print "Starting relaxation step"
relax_start = time.time()
if offload_enabled:
op.relax(2, offl_b.array, offl_a.array, 4, 2.0 / 3.0)
else:
op.relax(2, b, a, 4, 2.0 / 3.0)
relax_end = time.time()
print "--------------------------------------------------------------"
class SternheimerOperator:
"""Class implementing the linear operator in the Sternheimer equation.
The main purpose of this class is to implement the multiplication of the
linear operator (H - eps_nk) with a vector in the ``apply`` member
function. An additional ``matvec`` member function has been defined so that
instances of this class can be passed as a linear operator to scipy's
iterative Krylov solvers in ``scipy.sparse.linalg``.
"""
def __init__(self, hamiltonian, wfs, gd, dtype=float):
"""Save useful objects for the Sternheimer operator.
Parameters
----------
hamiltonian: Hamiltonian
Hamiltonian for a ground-state calculation.
wfs: WaveFunctions
Ground-state wave-functions.
gd: GridDescriptor
Grid on which the operator is defined.
dtype: dtype
dtype of the wave-function being operated on.
"""
self.hamiltonian = hamiltonian
self.kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
self.kpt_u = wfs.kpt_u
self.pt = wfs.pt
self.gd = gd
# Variables for k-point and band index
self.k = None
self.n = None
self.kplusq = None
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N, N)
self.dtype = dtype
def set_k(self, k):
"""Set k-vector for the Bloch-state in consideration."""
self.k = k
def set_band(self, n):
"""Set band index for the Bloch-state in consideration.
Note that n different from None has implications for the interpretation
of the input vector in the ``apply`` member function.
"""
self.n = n
def set_kplusq(self, kplusq):
"""Set q-vector of the perturbing potential.
Parameters
----------
kplusq: int
Index of the k+q vector.
"""
self.kplusq = kplusq
def apply(self, x_nG, y_nG):
"""Apply the linear operator in the Sternheimer equation to a vector.
For the eigenvalue term the k-point is the one of the state.
For the other terms the k-point to be used is the one given by the k+q
phase of the first-order of the state. Only for q=0 do the two coincide.
Parameters
----------
x_nG: ndarray
Vector(s) to which the Sternheimer operator is applied.
y_nG: ndarray
Resulting vector(s).
"""
assert x_nG.ndim in (3, 4)
assert x_nG.shape == y_nG.shape
assert tuple(self.gd.n_c) == x_nG.shape[-3:]
assert self.k is not None
# k
kpt = self.kpt_u[self.k]
# k+q
kplusqpt = self.kpt_u[self.kplusq]
# Kintetic energy
# k+q
self.kin.apply(x_nG, y_nG, phase_cd=kplusqpt.phase_cd)
# Local part of effective potential - no phase !!
self.hamiltonian.apply_local_potential(x_nG, y_nG, kpt.s)
# Non-local part from projectors (coefficients can not be reused)
shape = x_nG.shape[:-3]
P_ani = self.pt.dict(shape=shape)
# k+q
self.pt.integrate(x_nG, P_ani, q=kplusqpt.k)
for a, P_ni in P_ani.items():
dH_ii = unpack(self.hamiltonian.dH_asp[a][kpt.s])
P_ani[a] = np.dot(P_ni, dH_ii)
# k+q
self.pt.add(y_nG, P_ani, q=kplusqpt.k)
# XXX Eigenvalue term
if self.n is not None:
# k
y_nG -= kpt.eps_n[self.n] * x_nG
else:
for n, x_G in enumerate(x_nG):
# k
y_nG[n] -= kpt.eps_n[n] * x_G
# Project out undesired (numerical) components
# k+q
self.project(y_nG)
def project(self, x_nG):
"""Project the vector onto the unoccupied states at k+q.
The projection operator is defined as follows::
-- --
P = > |Psi ><Psi | = 1 - > |Psi ><Psi |
c,k+q -- c c -- v v
c k+q k+q v k+q k+q
"""
# It might be a good idea to move this functionality to its own class
assert x_nG.ndim == 4
assert self.kplusq is not None
# k+q
kplusqpt = self.kpt_u[self.kplusq]
# Occupied wave function
psit_nG = kplusqpt.psit_nG
# Project out occupied sub-space using blas routine gemm
m = len(x_nG)
n = len(psit_nG)
proj_mn = np.zeros((m, n), dtype=self.dtype)
gemm(self.gd.dv, psit_nG, x_nG, 0.0, proj_mn, 'c')
gemm(-1.0, psit_nG, proj_mn, 1.0, x_nG)
# Project out one orbital at a time
## for n, psit_G in enumerate(psit_nG):
##
## proj = self.gd.integrate(psit_G.conjugate() * x_nG)
## x_nG -= proj * psit_G
def matvec(self, x):
"""Matrix-vector multiplication for scipy's Krylov solvers.
This is a wrapper around the ``apply`` member function above. It allows
to multiply the sternheimer operator onto the 1-dimensional scipy
representation of a gpaw grid vector(s).
Parameters
----------
a: ndarray
1-dimensional array holding the representation of a gpaw grid
vector (possibly a set of vectors).
"""
# Check that a is 1-dimensional
assert len(x.shape) == 1
# Find the number of states in x
grid_shape = tuple(self.gd.n_c)
assert ((x.size % np.prod(grid_shape)) == 0), ("Incompatible array " +
"shapes")
# Number of states in the vector
N = x.size / np.prod(grid_shape)
# Output array
y_nG = self.gd.zeros(n=N, dtype=self.dtype)
shape = y_nG.shape
assert x.size == np.prod(y_nG.shape)
x_nG = x.reshape(shape)
self.apply(x_nG, y_nG)
y = y_nG.ravel()
return y
class FDWaveFunctions(FDPWWaveFunctions):
def __init__(self,
stencil,
diagksl,
orthoksl,
initksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, dtype, world, kd,
timer)
self.wd = self.gd # wave function descriptor
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype, allocate=False)
self.matrixoperator = MatrixOperator(orthoksl)
def set_setups(self, setups):
self.pt = LFC(self.gd, [setup.pt_j for setup in setups],
self.kpt_comm,
dtype=self.dtype,
forces=True)
FDPWWaveFunctions.set_setups(self, setups)
def set_positions(self, spos_ac):
if not self.kin.is_allocated():
self.kin.allocate()
FDPWWaveFunctions.set_positions(self, spos_ac)
def summary(self, fd):
fd.write('Mode: Finite-difference\n')
def make_preconditioner(self, block=1):
return Preconditioner(self.gd, self.kin, self.dtype, block)
def apply_pseudo_hamiltonian(self, kpt, hamiltonian, psit_xG, Htpsit_xG):
self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd)
hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s)
def add_orbital_density(self, nt_G, kpt, n):
if self.dtype == float:
axpy(1.0, kpt.psit_nG[n]**2, nt_G)
else:
axpy(1.0, kpt.psit_nG[n].real**2, nt_G)
axpy(1.0, kpt.psit_nG[n].imag**2, nt_G)
def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n):
# Used in calculation of response part of GLLB-potential
nt_G = nt_sG[kpt.s]
if self.dtype == float:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G**2, nt_G)
else:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G.real**2, nt_G)
axpy(f, psit_G.imag**2, nt_G)
# Hack used in delta-scf calculations:
if hasattr(kpt, 'c_on'):
assert self.bd.comm.size == 1
d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands),
dtype=complex)
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn += ne * np.outer(c_n.conj(), c_n)
for d_n, psi0_G in zip(d_nn, kpt.psit_nG):
for d, psi_G in zip(d_n, kpt.psit_nG):
if abs(d) > 1.e-12:
nt_G += (psi0_G.conj() * d * psi_G).real
def calculate_kinetic_energy_density(self, tauct, grad_v):
assert not hasattr(self.kpt_u[0], 'c_on')
if isinstance(self.kpt_u[0].psit_nG, TarFileReference):
raise RuntimeError('Wavefunctions have not been initialized.')
taut_sG = self.gd.zeros(self.nspins)
dpsit_G = self.gd.empty(dtype=self.dtype)
for kpt in self.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
for v in range(3):
grad_v[v](psit_G, dpsit_G, kpt.phase_cd)
axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s])
self.kpt_comm.sum(taut_sG)
self.band_comm.sum(taut_sG)
return taut_sG
def calculate_forces(self, hamiltonian, F_av):
# Calculate force-contribution from k-points:
F_av.fill(0.0)
F_aniv = self.pt.dict(self.bd.mynbands, derivative=True)
for kpt in self.kpt_u:
self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q)
for a, F_niv in F_aniv.items():
F_niv = F_niv.conj()
F_niv *= kpt.f_n[:, np.newaxis, np.newaxis]
dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
P_ni = kpt.P_ani[a]
F_vii = np.dot(np.dot(F_niv.transpose(), P_ni), dH_ii)
F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis]
dO_ii = hamiltonian.setups[a].dO_ii
F_vii -= np.dot(np.dot(F_niv.transpose(), P_ni), dO_ii)
F_av[a] += 2 * F_vii.real.trace(0, 1, 2)
# Hack used in delta-scf calculations:
if hasattr(kpt, 'c_on'):
assert self.bd.comm.size == 1
self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q) #XXX again
d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands),
dtype=complex)
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn += ne * np.outer(c_n.conj(), c_n)
for a, F_niv in F_aniv.items():
F_niv = F_niv.conj()
dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
Q_ni = np.dot(d_nn, kpt.P_ani[a])
F_vii = np.dot(np.dot(F_niv.transpose(), Q_ni), dH_ii)
F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis]
dO_ii = hamiltonian.setups[a].dO_ii
F_vii -= np.dot(np.dot(F_niv.transpose(), Q_ni), dO_ii)
F_av[a] += 2 * F_vii.real.trace(0, 1, 2)
self.bd.comm.sum(F_av, 0)
if self.bd.comm.rank == 0:
self.kpt_comm.sum(F_av, 0)
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
self.kin.estimate_memory(mem.subnode('Kinetic operator'))
class FDWaveFunctions(FDPWWaveFunctions):
mode = 'fd'
def __init__(self,
stencil,
diagksl,
orthoksl,
initksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
kptband_comm,
timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, dtype, world, kd,
kptband_comm, timer)
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype)
self.matrixoperator = MatrixOperator(self.orthoksl)
self.taugrad_v = None # initialized by MGGA functional
def empty(self, n=(), global_array=False, realspace=False, q=-1):
return self.gd.empty(n, self.dtype, global_array)
def integrate(self, a_xg, b_yg=None, global_integral=True):
return self.gd.integrate(a_xg, b_yg, global_integral)
def bytes_per_wave_function(self):
return self.gd.bytecount(self.dtype)
def set_setups(self, setups):
self.pt = LFC(self.gd, [setup.pt_j for setup in setups],
self.kd,
dtype=self.dtype,
forces=True)
FDPWWaveFunctions.set_setups(self, setups)
def set_positions(self, spos_ac):
FDPWWaveFunctions.set_positions(self, spos_ac)
def summary(self, fd):
fd.write('Wave functions: Uniform real-space grid\n')
fd.write('Kinetic energy operator: %s\n' % self.kin.description)
def make_preconditioner(self, block=1):
return Preconditioner(self.gd, self.kin, self.dtype, block)
def apply_pseudo_hamiltonian(self, kpt, hamiltonian, psit_xG, Htpsit_xG):
self.timer.start('Apply hamiltonian')
self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd)
hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s)
self.timer.stop('Apply hamiltonian')
def add_orbital_density(self, nt_G, kpt, n):
if self.dtype == float:
axpy(1.0, kpt.psit_nG[n]**2, nt_G)
else:
axpy(1.0, kpt.psit_nG[n].real**2, nt_G)
axpy(1.0, kpt.psit_nG[n].imag**2, nt_G)
def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n):
# Used in calculation of response part of GLLB-potential
nt_G = nt_sG[kpt.s]
if self.dtype == float:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G**2, nt_G)
else:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G.real**2, nt_G)
axpy(f, psit_G.imag**2, nt_G)
# Hack used in delta-scf calculations:
if hasattr(kpt, 'c_on'):
assert self.bd.comm.size == 1
d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands),
dtype=complex)
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn += ne * np.outer(c_n.conj(), c_n)
for d_n, psi0_G in zip(d_nn, kpt.psit_nG):
for d, psi_G in zip(d_n, kpt.psit_nG):
if abs(d) > 1.e-12:
nt_G += (psi0_G.conj() * d * psi_G).real
def calculate_kinetic_energy_density(self):
if self.taugrad_v is None:
self.taugrad_v = [
Gradient(self.gd, v, n=3, dtype=self.dtype).apply
for v in range(3)
]
assert not hasattr(self.kpt_u[0], 'c_on')
if self.kpt_u[0].psit_nG is None:
raise RuntimeError('No wavefunctions yet')
if isinstance(self.kpt_u[0].psit_nG, FileReference):
# XXX initialize
raise RuntimeError('Wavefunctions have not been initialized.')
taut_sG = self.gd.zeros(self.nspins)
dpsit_G = self.gd.empty(dtype=self.dtype)
for kpt in self.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
for v in range(3):
self.taugrad_v[v](psit_G, dpsit_G, kpt.phase_cd)
axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s])
self.kd.comm.sum(taut_sG)
self.band_comm.sum(taut_sG)
return taut_sG
def apply_mgga_orbital_dependent_hamiltonian(self, kpt, psit_xG, Htpsit_xG,
dH_asp, dedtaut_G):
a_G = self.gd.empty(dtype=psit_xG.dtype)
for psit_G, Htpsit_G in zip(psit_xG, Htpsit_xG):
for v in range(3):
self.taugrad_v[v](psit_G, a_G, kpt.phase_cd)
self.taugrad_v[v](dedtaut_G * a_G, a_G, kpt.phase_cd)
axpy(-0.5, a_G, Htpsit_G)
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 write(self, writer, write_wave_functions=False):
writer['Mode'] = 'fd'
if not write_wave_functions:
return
writer.add(
'PseudoWaveFunctions',
('nspins', 'nibzkpts', 'nbands', 'ngptsx', 'ngptsy', 'ngptsz'),
dtype=self.dtype)
if hasattr(writer, 'hdf5'):
parallel = (self.world.size > 1)
for kpt in self.kpt_u:
indices = [kpt.s, kpt.k]
indices.append(self.bd.get_slice())
indices += self.gd.get_slice()
writer.fill(kpt.psit_nG, parallel=parallel, *indices)
else:
for s in range(self.nspins):
for k in range(self.kd.nibzkpts):
for n in range(self.bd.nbands):
psit_G = self.get_wave_function_array(n, k, s)
writer.fill(psit_G, s, k, n)
def read(self, reader, hdf5):
if ((not hdf5 and self.bd.comm.size == 1)
or (hdf5 and self.world.size == 1)):
# We may not be able to keep all the wave
# functions in memory - so psit_nG will be a special type of
# array that is really just a reference to a file:
for kpt in self.kpt_u:
kpt.psit_nG = reader.get_reference('PseudoWaveFunctions',
(kpt.s, kpt.k))
else:
for kpt in self.kpt_u:
kpt.psit_nG = self.empty(self.bd.mynbands)
if hdf5:
indices = [kpt.s, kpt.k]
indices.append(self.bd.get_slice())
indices += self.gd.get_slice()
reader.get('PseudoWaveFunctions',
out=kpt.psit_nG,
parallel=(self.world.size > 1),
*indices)
else:
# Read band by band to save memory
for myn, psit_G in enumerate(kpt.psit_nG):
n = self.bd.global_index(myn)
if self.gd.comm.rank == 0:
big_psit_G = np.array(
reader.get('PseudoWaveFunctions', kpt.s, kpt.k,
n), self.dtype)
else:
big_psit_G = None
self.gd.distribute(big_psit_G, psit_G)
def initialize_from_lcao_coefficients(self, basis_functions, mynbands):
for kpt in self.kpt_u:
kpt.psit_nG = self.gd.zeros(self.bd.mynbands, self.dtype)
basis_functions.lcao_to_grid(kpt.C_nM, kpt.psit_nG[:mynbands],
kpt.q)
kpt.C_nM = None
if use_mic:
kpt.psit_nG_mic = stream.bind(kpt.psit_nG)
stream.sync()
def random_wave_functions(self, nao):
"""Generate random wave functions."""
gpts = self.gd.N_c[0] * self.gd.N_c[1] * self.gd.N_c[2]
if self.bd.nbands < gpts / 64:
gd1 = self.gd.coarsen()
gd2 = gd1.coarsen()
psit_G1 = gd1.empty(dtype=self.dtype)
psit_G2 = gd2.empty(dtype=self.dtype)
interpolate2 = Transformer(gd2, gd1, 1, self.dtype).apply
interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply
shape = tuple(gd2.n_c)
scale = np.sqrt(12 / abs(np.linalg.det(gd2.cell_cv)))
old_state = np.random.get_state()
np.random.seed(4 + self.world.rank)
for kpt in self.kpt_u:
for psit_G in kpt.psit_nG[nao:]:
if self.dtype == float:
psit_G2[:] = (np.random.random(shape) - 0.5) * scale
else:
psit_G2.real = (np.random.random(shape) - 0.5) * scale
psit_G2.imag = (np.random.random(shape) - 0.5) * scale
interpolate2(psit_G2, psit_G1, kpt.phase_cd)
interpolate1(psit_G1, psit_G, kpt.phase_cd)
np.random.set_state(old_state)
elif gpts / 64 <= self.bd.nbands < gpts / 8:
gd1 = self.gd.coarsen()
psit_G1 = gd1.empty(dtype=self.dtype)
interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply
shape = tuple(gd1.n_c)
scale = np.sqrt(12 / abs(np.linalg.det(gd1.cell_cv)))
old_state = np.random.get_state()
np.random.seed(4 + self.world.rank)
for kpt in self.kpt_u:
for psit_G in kpt.psit_nG[nao:]:
if self.dtype == float:
psit_G1[:] = (np.random.random(shape) - 0.5) * scale
else:
psit_G1.real = (np.random.random(shape) - 0.5) * scale
psit_G1.imag = (np.random.random(shape) - 0.5) * scale
interpolate1(psit_G1, psit_G, kpt.phase_cd)
np.random.set_state(old_state)
else:
shape = tuple(self.gd.n_c)
scale = np.sqrt(12 / abs(np.linalg.det(self.gd.cell_cv)))
old_state = np.random.get_state()
np.random.seed(4 + self.world.rank)
for kpt in self.kpt_u:
for psit_G in kpt.psit_nG[nao:]:
if self.dtype == float:
psit_G[:] = (np.random.random(shape) - 0.5) * scale
else:
psit_G.real = (np.random.random(shape) - 0.5) * scale
psit_G.imag = (np.random.random(shape) - 0.5) * scale
np.random.set_state(old_state)
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
class FDWaveFunctions(FDPWWaveFunctions):
def __init__(self, stencil, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer)
self.wd = self.gd # wave function descriptor
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype, allocate=False)
self.matrixoperator = MatrixOperator(orthoksl)
def set_setups(self, setups):
self.pt = LFC(self.gd, [setup.pt_j for setup in setups],
self.kpt_comm, dtype=self.dtype, forces=True)
FDPWWaveFunctions.set_setups(self, setups)
def set_positions(self, spos_ac):
if not self.kin.is_allocated():
self.kin.allocate()
FDPWWaveFunctions.set_positions(self, spos_ac)
def summary(self, fd):
fd.write('Mode: Finite-difference\n')
def make_preconditioner(self, block=1):
return Preconditioner(self.gd, self.kin, self.dtype, block)
def apply_pseudo_hamiltonian(self, kpt, hamiltonian, psit_xG, Htpsit_xG):
self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd)
hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s)
def add_orbital_density(self, nt_G, kpt, n):
if self.dtype == float:
axpy(1.0, kpt.psit_nG[n]**2, nt_G)
else:
axpy(1.0, kpt.psit_nG[n].real**2, nt_G)
axpy(1.0, kpt.psit_nG[n].imag**2, nt_G)
def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n):
# Used in calculation of response part of GLLB-potential
nt_G = nt_sG[kpt.s]
if self.dtype == float:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G**2, nt_G)
else:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G.real**2, nt_G)
axpy(f, psit_G.imag**2, nt_G)
# Hack used in delta-scf calculations:
if hasattr(kpt, 'c_on'):
assert self.bd.comm.size == 1
d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands),
dtype=complex)
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn += ne * np.outer(c_n.conj(), c_n)
for d_n, psi0_G in zip(d_nn, kpt.psit_nG):
for d, psi_G in zip(d_n, kpt.psit_nG):
if abs(d) > 1.e-12:
nt_G += (psi0_G.conj() * d * psi_G).real
def calculate_kinetic_energy_density(self, tauct, grad_v):
assert not hasattr(self.kpt_u[0], 'c_on')
if isinstance(self.kpt_u[0].psit_nG, TarFileReference):
raise RuntimeError('Wavefunctions have not been initialized.')
taut_sG = self.gd.zeros(self.nspins)
dpsit_G = self.gd.empty(dtype=self.dtype)
for kpt in self.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
for v in range(3):
grad_v[v](psit_G, dpsit_G, kpt.phase_cd)
axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s])
self.kpt_comm.sum(taut_sG)
self.band_comm.sum(taut_sG)
return taut_sG
def calculate_forces(self, hamiltonian, F_av):
# Calculate force-contribution from k-points:
F_av.fill(0.0)
F_aniv = self.pt.dict(self.bd.mynbands, derivative=True)
for kpt in self.kpt_u:
self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q)
for a, F_niv in F_aniv.items():
F_niv = F_niv.conj()
F_niv *= kpt.f_n[:, np.newaxis, np.newaxis]
dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
P_ni = kpt.P_ani[a]
F_vii = np.dot(np.dot(F_niv.transpose(), P_ni), dH_ii)
F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis]
dO_ii = hamiltonian.setups[a].dO_ii
F_vii -= np.dot(np.dot(F_niv.transpose(), P_ni), dO_ii)
F_av[a] += 2 * F_vii.real.trace(0, 1, 2)
# Hack used in delta-scf calculations:
if hasattr(kpt, 'c_on'):
assert self.bd.comm.size == 1
self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q) #XXX again
d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands),
dtype=complex)
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn += ne * np.outer(c_n.conj(), c_n)
for a, F_niv in F_aniv.items():
F_niv = F_niv.conj()
dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s])
Q_ni = np.dot(d_nn, kpt.P_ani[a])
F_vii = np.dot(np.dot(F_niv.transpose(), Q_ni), dH_ii)
F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis]
dO_ii = hamiltonian.setups[a].dO_ii
F_vii -= np.dot(np.dot(F_niv.transpose(), Q_ni), dO_ii)
F_av[a] += 2 * F_vii.real.trace(0, 1, 2)
self.bd.comm.sum(F_av, 0)
if self.bd.comm.rank == 0:
self.kpt_comm.sum(F_av, 0)
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
self.kin.estimate_memory(mem.subnode('Kinetic operator'))
class FDWaveFunctions(FDPWWaveFunctions):
def __init__(self, stencil, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
dtype, world, kd, timer)
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype)
self.matrixoperator = MatrixOperator(self.orthoksl)
self.taugrad_v = None # initialized by MGGA functional
def empty(self, n=(), global_array=False, realspace=False, q=-1):
return self.gd.empty(n, self.dtype, global_array)
def integrate(self, a_xg, b_yg=None, global_integral=True):
return self.gd.integrate(a_xg, b_yg, global_integral)
def bytes_per_wave_function(self):
return self.gd.bytecount(self.dtype)
def set_setups(self, setups):
self.pt = LFC(self.gd, [setup.pt_j for setup in setups],
self.kd, dtype=self.dtype, forces=True)
FDPWWaveFunctions.set_setups(self, setups)
def set_positions(self, spos_ac):
FDPWWaveFunctions.set_positions(self, spos_ac)
def summary(self, fd):
fd.write('Wave functions: Uniform real-space grid\n')
fd.write('Kinetic energy operator: %s\n' % self.kin.description)
def make_preconditioner(self, block=1):
return Preconditioner(self.gd, self.kin, self.dtype, block)
def apply_pseudo_hamiltonian(self, kpt, hamiltonian, psit_xG, Htpsit_xG):
self.timer.start('Apply hamiltonian')
self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd)
hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s)
self.timer.stop('Apply hamiltonian')
def add_orbital_density(self, nt_G, kpt, n):
if self.dtype == float:
axpy(1.0, kpt.psit_nG[n]**2, nt_G)
else:
axpy(1.0, kpt.psit_nG[n].real**2, nt_G)
axpy(1.0, kpt.psit_nG[n].imag**2, nt_G)
def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n):
# Used in calculation of response part of GLLB-potential
nt_G = nt_sG[kpt.s]
if self.dtype == float:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G**2, nt_G)
else:
for f, psit_G in zip(f_n, kpt.psit_nG):
axpy(f, psit_G.real**2, nt_G)
axpy(f, psit_G.imag**2, nt_G)
# Hack used in delta-scf calculations:
if hasattr(kpt, 'c_on'):
assert self.bd.comm.size == 1
d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands),
dtype=complex)
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn += ne * np.outer(c_n.conj(), c_n)
for d_n, psi0_G in zip(d_nn, kpt.psit_nG):
for d, psi_G in zip(d_n, kpt.psit_nG):
if abs(d) > 1.e-12:
nt_G += (psi0_G.conj() * d * psi_G).real
def calculate_kinetic_energy_density(self):
if self.taugrad_v is None:
self.taugrad_v = [
Gradient(self.gd, v, n=3, dtype=self.dtype).apply
for v in range(3)]
assert not hasattr(self.kpt_u[0], 'c_on')
if self.kpt_u[0].psit_nG is None:
raise RuntimeError('No wavefunctions yet')
if isinstance(self.kpt_u[0].psit_nG, FileReference):
# XXX initialize
raise RuntimeError('Wavefunctions have not been initialized.')
taut_sG = self.gd.zeros(self.nspins)
dpsit_G = self.gd.empty(dtype=self.dtype)
for kpt in self.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
for v in range(3):
self.taugrad_v[v](psit_G, dpsit_G, kpt.phase_cd)
axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s])
self.kpt_comm.sum(taut_sG)
self.band_comm.sum(taut_sG)
return taut_sG
def apply_mgga_orbital_dependent_hamiltonian(self, kpt, psit_xG,
Htpsit_xG, dH_asp,
dedtaut_G):
a_G = self.gd.empty(dtype=psit_xG.dtype)
for psit_G, Htpsit_G in zip(psit_xG, Htpsit_xG):
for v in range(3):
self.taugrad_v[v](psit_G, a_G, kpt.phase_cd)
self.taugrad_v[v](dedtaut_G * a_G, a_G, kpt.phase_cd)
axpy(-0.5, a_G, Htpsit_G)
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 write(self, writer, write_wave_functions=False):
writer['Mode'] = 'fd'
if not write_wave_functions:
return
writer.add('PseudoWaveFunctions',
('nspins', 'nibzkpts', 'nbands',
'ngptsx', 'ngptsy', 'ngptsz'),
dtype=self.dtype)
if hasattr(writer, 'hdf5'):
parallel = (self.world.size > 1)
for kpt in self.kpt_u:
indices = [kpt.s, kpt.k]
indices.append(self.bd.get_slice())
indices += self.gd.get_slice()
writer.fill(kpt.psit_nG, parallel=parallel, *indices)
else:
for s in range(self.nspins):
for k in range(self.nibzkpts):
for n in range(self.bd.nbands):
psit_G = self.get_wave_function_array(n, k, s)
writer.fill(psit_G, s, k, n)
def read(self, reader, hdf5):
if ((not hdf5 and self.bd.comm.size == 1) or
(hdf5 and self.world.size == 1)):
# We may not be able to keep all the wave
# functions in memory - so psit_nG will be a special type of
# array that is really just a reference to a file:
for kpt in self.kpt_u:
kpt.psit_nG = reader.get_reference('PseudoWaveFunctions',
(kpt.s, kpt.k))
else:
for kpt in self.kpt_u:
kpt.psit_nG = self.empty(self.bd.mynbands)
if hdf5:
indices = [kpt.s, kpt.k]
indices.append(self.bd.get_slice())
indices += self.gd.get_slice()
reader.get('PseudoWaveFunctions', out=kpt.psit_nG,
parallel=(self.world.size > 1), *indices)
else:
# Read band by band to save memory
for myn, psit_G in enumerate(kpt.psit_nG):
n = self.bd.global_index(myn)
if self.gd.comm.rank == 0:
big_psit_G = np.array(
reader.get('PseudoWaveFunctions',
kpt.s, kpt.k, n),
self.dtype)
else:
big_psit_G = None
self.gd.distribute(big_psit_G, psit_G)
def initialize_from_lcao_coefficients(self, basis_functions, mynbands):
for kpt in self.kpt_u:
kpt.psit_nG = self.gd.zeros(self.bd.mynbands, self.dtype)
basis_functions.lcao_to_grid(kpt.C_nM,
kpt.psit_nG[:mynbands], kpt.q)
kpt.C_nM = None
def random_wave_functions(self, nao):
"""Generate random wave functions."""
gpts = self.gd.N_c[0] * self.gd.N_c[1] * self.gd.N_c[2]
if self.bd.nbands < gpts / 64:
gd1 = self.gd.coarsen()
gd2 = gd1.coarsen()
psit_G1 = gd1.empty(dtype=self.dtype)
psit_G2 = gd2.empty(dtype=self.dtype)
interpolate2 = Transformer(gd2, gd1, 1, self.dtype).apply
interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply
shape = tuple(gd2.n_c)
scale = np.sqrt(12 / abs(np.linalg.det(gd2.cell_cv)))
old_state = np.random.get_state()
np.random.seed(4 + self.world.rank)
for kpt in self.kpt_u:
for psit_G in kpt.psit_nG[nao:]:
if self.dtype == float:
psit_G2[:] = (np.random.random(shape) - 0.5) * scale
else:
psit_G2.real = (np.random.random(shape) - 0.5) * scale
psit_G2.imag = (np.random.random(shape) - 0.5) * scale
interpolate2(psit_G2, psit_G1, kpt.phase_cd)
interpolate1(psit_G1, psit_G, kpt.phase_cd)
np.random.set_state(old_state)
elif gpts / 64 <= self.bd.nbands < gpts / 8:
gd1 = self.gd.coarsen()
psit_G1 = gd1.empty(dtype=self.dtype)
interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply
shape = tuple(gd1.n_c)
scale = np.sqrt(12 / abs(np.linalg.det(gd1.cell_cv)))
old_state = np.random.get_state()
np.random.seed(4 + self.world.rank)
for kpt in self.kpt_u:
for psit_G in kpt.psit_nG[nao:]:
if self.dtype == float:
psit_G1[:] = (np.random.random(shape) - 0.5) * scale
else:
psit_G1.real = (np.random.random(shape) - 0.5) * scale
psit_G1.imag = (np.random.random(shape) - 0.5) * scale
interpolate1(psit_G1, psit_G, kpt.phase_cd)
np.random.set_state(old_state)
else:
shape = tuple(self.gd.n_c)
scale = np.sqrt(12 / abs(np.linalg.det(self.gd.cell_cv)))
old_state = np.random.get_state()
np.random.seed(4 + self.world.rank)
for kpt in self.kpt_u:
for psit_G in kpt.psit_nG[nao:]:
if self.dtype == float:
psit_G[:] = (np.random.random(shape) - 0.5) * scale
else:
psit_G.real = (np.random.random(shape) - 0.5) * scale
psit_G.imag = (np.random.random(shape) - 0.5) * scale
np.random.set_state(old_state)
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
def set_grid_descriptor(self, gd):
self.gd = gd
axes = np.arange(3)
pbc_c = np.array(gd.pbc_c, dtype=bool)
periodic_axes = axes[pbc_c]
non_periodic_axes = axes[np.logical_not(pbc_c)]
# Find out which axes are orthogonal (0, 1 or 3)
# Note that one expects that the axes are always rotated in
# conventional form, thus for all axes to be
# classified as orthogonal, the cell_cv needs to be diagonal.
# This may always be achieved by rotating
# the unit-cell along with the atoms. The classification is
# inherited from grid_descriptor.orthogonal.
dotprods = np.dot(gd.cell_cv, gd.cell_cv.T)
# For each direction, check whether there is only one nonzero
# element in that row (necessarily being the diagonal element,
# since this is a cell vector length and must be > 0).
orthogonal_c = (np.abs(dotprods) > 1e-10).sum(axis=0) == 1
assert sum(orthogonal_c) in [0, 1, 3]
non_orthogonal_axes = axes[np.logical_not(orthogonal_c)]
if not all(pbc_c | orthogonal_c):
raise BadAxesError('Each axis must be periodic or orthogonal '
'to other axes. But we have pbc={} '
'and orthogonal={}'.format(
pbc_c.astype(int),
orthogonal_c.astype(int)))
# We sort them, and pick the longest non-periodic axes as the
# cholesky axis.
sorted_non_periodic_axes = sorted(non_periodic_axes,
key=lambda c: gd.N_c[c])
if self.use_cholesky:
if len(sorted_non_periodic_axes) > 0:
cholesky_axes = [sorted_non_periodic_axes[-1]]
if cholesky_axes[0] in non_orthogonal_axes:
msg = ('Cholesky axis cannot be non-orthogonal. '
'Do you really want a non-orthogonal non-periodic '
'axis? If so, run with use_cholesky=False.')
raise NotImplementedError(msg)
fst_axes = sorted_non_periodic_axes[0:-1]
else:
cholesky_axes = []
fst_axes = []
else:
cholesky_axes = []
fst_axes = sorted_non_periodic_axes
fft_axes = list(periodic_axes)
(self.cholesky_axes, self.fst_axes,
self.fft_axes) = cholesky_axes, fst_axes, fft_axes
fftfst_axes = self.fft_axes + self.fst_axes
axes = self.fft_axes + self.fst_axes + self.cholesky_axes
self.axes = axes
gd_x = [self.gd]
# Create xy flat decomposition (where x=axes[0] and y=axes[1])
domain = gd.N_c.copy()
domain[axes[0]] = 1
domain[axes[1]] = 1
parsize_c = decompose_domain(domain, gd.comm.size)
gd_x.append(gd.new_descriptor(parsize_c=parsize_c))
# Create z flat decomposition
domain = gd.N_c.copy()
domain[axes[2]] = 1
parsize_c = decompose_domain(domain, gd.comm.size)
gd_x.append(gd.new_descriptor(parsize_c=parsize_c))
self.gd_x = gd_x
# Calculate eigenvalues in fst/fft decomposition for
# non-cholesky axes in parallel
r_cx = np.indices(gd_x[-1].n_c)
r_cx += gd_x[-1].beg_c[:, np.newaxis, np.newaxis, np.newaxis]
r_cx = r_cx.astype(complex)
for c, axis in enumerate(fftfst_axes):
r_cx[axis] *= 2j * np.pi / gd_x[-1].N_c[axis]
if axis in fst_axes:
r_cx[axis] /= 2
for c, axis in enumerate(cholesky_axes):
r_cx[axis] = 0.0
np.exp(r_cx, out=r_cx)
fft_lambdas = np.zeros_like(r_cx[0], dtype=complex)
laplace = Laplace(gd_x[-1], -0.25 / pi, self.nn)
self.stencil_description = laplace.description
for coeff, offset_c in zip(laplace.coef_p, laplace.offset_pc):
offset_c = np.array(offset_c)
if not any(offset_c):
# The centerpoint is handled with (temp-1.0)
continue
non_zero_axes, = np.where(offset_c)
if set(non_zero_axes).issubset(fftfst_axes):
temp = np.ones_like(fft_lambdas)
for c, axis in enumerate(fftfst_axes):
temp *= r_cx[axis]**offset_c[axis]
fft_lambdas += coeff * (temp - 1.0)
assert np.linalg.norm(fft_lambdas.imag) < 1e-10
fft_lambdas = fft_lambdas.real.copy() # arr.real is not contiguous
# If there is no Cholesky decomposition, the system is already
# fully diagonal and we can directly invert the linear problem
# by dividing with the eigenvalues.
assert len(cholesky_axes) == 0
# if len(cholesky_axes) == 0:
with np.errstate(divide='ignore'):
self.inv_fft_lambdas = np.where(
np.abs(fft_lambdas) > 1e-10, 1.0 / fft_lambdas, 0)
""" Test the finite difference stencil """ from gpaw.grid_descriptor import GridDescriptor from gpaw.fd_operators import Laplace from gpaw.test import equal import numpy as np gpts = 64 nbands = 120 gd = GridDescriptor([gpts, gpts, gpts]) a = gd.empty(nbands) b = gd.empty(nbands) np.random.seed(10) a[:] = np.random.random(a.shape) op = Laplace(gd, 1.0, 3).apply op(a, b) equal(np.sum(b), -7.43198208966e-05, 1e-9)
import numpy as np
from gpaw.grid_descriptor import GridDescriptor
from gpaw.fd_operators import Laplace
from gpaw.mpi import rank, size, world
G = 40
N = 2 * 3 * 2 * 5 * 7 * 8 * 3 * 11
B = 2
h = 0.2
a = h * R
M = N // B
assert M * B == N
D = size // B
assert D * B == size
r = rank // B * B
domain_comm = mpi.world.new_communicator(np.arange(r, r + D))
band_comm = mpi.world.new_communicator(np.arange(rank, size, D))
gd = GridDescriptor((G, G, G), (a, a, a), comm=domain_comm)
np.random.seed(rank)
psit_mG = np.random.uniform(size=(M, ) + tuple(gd.n_c))
send_mG = gd.empty(M)
recv_mG = gd.empty(M)
laplace = [Laplace(g, n=1).apply for g in gd]
for i in range(B // 2):
rrequest = band_comm.reveive(recv_mG, (rank + D) % size, 42, 0)
srequest = band_comm.send(send_mG, (rank - D) % size, 17, 0)
# Strange one!
continue
print('------------------')
print(name, D)
print(cell[0])
print(cell[1])
print(cell[2])
for n in range(1, 6):
N = 2 * n + 2
gd = GridDescriptor((N, N, N), cell)
b_g = gd.zeros()
r_gv = gd.get_grid_point_coordinates().transpose((1, 2, 3, 0))
c_v = gd.cell_cv.sum(0) / 2
r_gv -= c_v
lap = Laplace(gd, n=n)
grad_v = [Gradient(gd, v, n=n) for v in range(3)]
assert lap.npoints == D * 2 * n + 1
for m in range(0, 2 * n + 1):
for ix in range(m + 1):
for iy in range(m - ix + 1):
iz = m - ix - iy
a_g = (r_gv**(ix, iy, iz)).prod(3)
if ix + iy + iz == 2 and max(ix, iy, iz) == 2:
r = 2.0
else:
r = 0.0
lap.apply(a_g, b_g)
e = b_g[n + 1, n + 1, n + 1] - r
assert abs(e) < 2e-12, e
for v in range(3):
def test(cellno, cellname, cell_cv, idiv, pbc, nn):
N_c = h2gpts(0.12, cell_cv, idiv=idiv)
if idiv == 1:
N_c += 1 - N_c % 2 # We want especially to test uneven grids
gd = GridDescriptor(N_c, cell_cv, pbc_c=pbc)
rho_g = gd.zeros()
phi_g = gd.zeros()
rho_g[:] = -0.3 + rng.rand(*rho_g.shape)
# Neutralize charge:
charge = gd.integrate(rho_g)
magic = gd.get_size_of_global_array().prod()
rho_g -= charge / gd.dv / magic
charge = gd.integrate(rho_g)
assert abs(charge) < 1e-12
# Check use_cholesky=True/False ?
from gpaw.poisson import FDPoissonSolver
ps = FastPoissonSolver(nn=nn)
#print('setgrid')
# Will raise BadAxesError for some pbc/cell combinations
ps.set_grid_descriptor(gd)
ps.solve(phi_g, rho_g)
laplace = Laplace(gd, scale=-1.0 / (4.0 * np.pi), n=nn)
def get_residual_err(phi_g):
rhotest_g = gd.zeros()
laplace.apply(phi_g, rhotest_g)
residual = np.abs(rhotest_g - rho_g)
# Residual is not accurate at end of non-periodic directions
# except for nn=1 (since effectively we use the right stencil
# only for nn=1 at the boundary).
#
# To do this check correctly, the Laplacian should have lower
# nn at the boundaries. Therefore we do not test the residual
# at these ends, only in between, by zeroing the bad ones:
if nn > 1:
exclude_points = nn - 1
for c in range(3):
if nn > 1 and not pbc[c]:
# get view ehere axis c refers becomes zeroth dimension:
X = residual.transpose(c, (c + 1) % 3, (c + 2) % 3)
if gd.beg_c[c] == 1:
X[:exclude_points] = 0.0
if gd.end_c[c] == gd.N_c[c]:
X[-exclude_points:] = 0.0
return residual.max()
maxerr = get_residual_err(phi_g)
pbcstring = '{}{}{}'.format(*pbc)
if 0:
ps2 = FDPoissonSolver(relax='J', nn=nn, eps=1e-18)
ps2.set_grid_descriptor(gd)
phi2_g = gd.zeros()
ps2.solve(phi2_g, rho_g)
phimaxerr = np.abs(phi2_g - phi_g).max()
maxerr2 = get_residual_err(phi2_g)
msg = ('{:2d} {:8s} pbc={} err={:8.5e} err[J]={:8.5e} '
'err[phi]={:8.5e} nn={:1d}'
.format(cellno, cellname, pbcstring, maxerr, maxerr2,
phimaxerr, nn))
state = 'ok' if maxerr < tolerance else 'FAIL'
msg = ('{:2d} {:8s} grid={} pbc={} err[fast]={:8.5e} nn={:1d} {}'
.format(cellno, cellname, N_c, pbcstring, maxerr, nn, state))
if world.rank == 0:
print(msg)
return maxerr
# Set up band and grid descriptors:
bd = BandDescriptor(N, band_comm, False)
gd = GridDescriptor((G, G, G), (a, a, a), True, domain_comm, parsize_c=D)
ksl = BandLayouts(gd, bd, block_comm, float)
# Random wave functions:
psit_mG = gd.empty(M)
for m in range(M):
np.random.seed(world.rank * M + m)
psit_mG[m] = np.random.uniform(-0.5, 0.5, tuple(gd.n_c))
if world.rank == 0:
print('Size of wave function array:', psit_mG.shape)
P_ani = {0: psit_mG[:, :2, 0, 0].copy(), 1: psit_mG[:, -1, -1, -3:].copy()}
kin = Laplace(gd, -0.5, 2).apply
vt_G = gd.empty()
vt_G.fill(0.567)
def run(psit_mG):
overlap = MatrixOperator(ksl, J)
def H(psit_xG):
Htpsit_xG = np.empty_like(psit_xG)
kin(psit_xG, Htpsit_xG)
for psit_G, y_G in zip(psit_xG, Htpsit_xG):
y_G += vt_G * psit_G
return Htpsit_xG
dH_aii = {0: np.ones((2, 2)) * 0.123, 1: np.ones((3, 3)) * 0.321}
class SternheimerOperator:
"""Class implementing the linear operator in the Sternheimer equation.
The main purpose of this class is to implement the multiplication of the
linear operator (H - eps_nk) with a vector in the ``apply`` member
function. An additional ``matvec`` member function has been defined so that
instances of this class can be passed as a linear operator to scipy's
iterative Krylov solvers in ``scipy.sparse.linalg``.
"""
def __init__(self, hamiltonian, wfs, gd, dtype=float):
"""Save useful objects for the Sternheimer operator.
Parameters
----------
hamiltonian: Hamiltonian
Hamiltonian for a ground-state calculation.
wfs: WaveFunctions
Ground-state wave-functions.
gd: GridDescriptor
Grid on which the operator is defined.
dtype: dtype
dtype of the wave-function being operated on.
"""
self.hamiltonian = hamiltonian
self.kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
self.kpt_u = wfs.kpt_u
self.pt = wfs.pt
self.gd = gd
# Variables for k-point and band index
self.k = None
self.n = None
self.kplusq = None
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N, N)
self.dtype = dtype
def set_k(self, k):
"""Set k-vector for the Bloch-state in consideration."""
self.k = k
def set_band(self, n):
"""Set band index for the Bloch-state in consideration.
Note that n different from None has implications for the interpretation
of the input vector in the ``apply`` member function.
"""
self.n = n
def set_kplusq(self, kplusq):
"""Set q-vector of the perturbing potential.
Parameters
----------
kplusq: int
Index of the k+q vector.
"""
self.kplusq = kplusq
def apply(self, x_nG, y_nG):
"""Apply the linear operator in the Sternheimer equation to a vector.
For the eigenvalue term the k-point is the one of the state.
For the other terms the k-point to be used is the one given by the k+q
phase of the first-order of the state. Only for q=0 do the two coincide.
Parameters
----------
x_nG: ndarray
Vector(s) to which the Sternheimer operator is applied.
y_nG: ndarray
Resulting vector(s).
"""
assert x_nG.ndim in (3, 4)
assert x_nG.shape == y_nG.shape
assert tuple(self.gd.n_c) == x_nG.shape[-3:]
assert self.k is not None
# k
kpt = self.kpt_u[self.k]
# k+q
kplusqpt = self.kpt_u[self.kplusq]
# Kintetic energy
# k+q
self.kin.apply(x_nG, y_nG, phase_cd=kplusqpt.phase_cd)
# Local part of effective potential - no phase !!
self.hamiltonian.apply_local_potential(x_nG, y_nG, kpt.s)
# Non-local part from projectors (coefficients can not be reused)
shape = x_nG.shape[:-3]
P_ani = self.pt.dict(shape=shape)
# k+q
self.pt.integrate(x_nG, P_ani, q=kplusqpt.k)
for a, P_ni in P_ani.items():
dH_ii = unpack(self.hamiltonian.dH_asp[a][kpt.s])
P_ani[a] = np.dot(P_ni, dH_ii)
# k+q
self.pt.add(y_nG, P_ani, q=kplusqpt.k)
# XXX Eigenvalue term
if self.n is not None:
# k
y_nG -= kpt.eps_n[self.n] * x_nG
else:
for n, x_G in enumerate(x_nG):
# k
y_nG[n] -= kpt.eps_n[n] * x_G
# Project out undesired (numerical) components
# k+q
self.project(y_nG)
def project(self, x_nG):
"""Project the vector onto the unoccupied states at k+q.
The projection operator is defined as follows::
-- --
P = > |Psi ><Psi | = 1 - > |Psi ><Psi |
c,k+q -- c c -- v v
c k+q k+q v k+q k+q
"""
# It might be a good idea to move this functionality to its own class
assert x_nG.ndim == 4
assert self.kplusq is not None
# k+q
kplusqpt = self.kpt_u[self.kplusq]
# Occupied wave function
psit_nG = kplusqpt.psit_nG
# Project out occupied sub-space using blas routine gemm
m = len(x_nG)
n = len(psit_nG)
proj_mn = np.zeros((m, n), dtype=self.dtype)
gemm(self.gd.dv, psit_nG, x_nG, 0.0, proj_mn, 'c')
gemm(-1.0, psit_nG, proj_mn, 1.0, x_nG)
# Project out one orbital at a time
## for n, psit_G in enumerate(psit_nG):
##
## proj = self.gd.integrate(psit_G.conjugate() * x_nG)
## x_nG -= proj * psit_G
def matvec(self, x):
"""Matrix-vector multiplication for scipy's Krylov solvers.
This is a wrapper around the ``apply`` member function above. It allows
to multiply the sternheimer operator onto the 1-dimensional scipy
representation of a gpaw grid vector(s).
Parameters
----------
a: ndarray
1-dimensional array holding the representation of a gpaw grid
vector (possibly a set of vectors).
"""
# Check that a is 1-dimensional
assert len(x.shape) == 1
# Find the number of states in x
grid_shape = tuple(self.gd.n_c)
assert ((x.size % np.prod(grid_shape)) == 0), ("Incompatible array " +
"shapes")
# Number of states in the vector
N = x.size / np.prod(grid_shape)
# Output array
y_nG = self.gd.zeros(n=N, dtype=self.dtype)
shape = y_nG.shape
assert x.size == np.prod(y_nG.shape)
x_nG = x.reshape(shape)
self.apply(x_nG, y_nG)
y = y_nG.ravel()
return y