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 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 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)
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
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
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)
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 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):
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 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):
mode = 'fd'
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 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, atom_partition=None):
FDPWWaveFunctions.set_positions(self, spos_ac, atom_partition)
def __str__(self):
s = 'Wave functions: Uniform real-space grid\n'
s += ' Kinetic energy operator: %s\n' % self.kin.description
return s + FDPWWaveFunctions.__str__(self)
def make_preconditioner(self, block=1):
return Preconditioner(self.gd, self.kin, self.dtype, block)
def apply_pseudo_hamiltonian(self, kpt, ham, psit_xG, Htpsit_xG):
self.timer.start('Apply hamiltonian')
self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd)
ham.apply_local_potential(psit_xG, Htpsit_xG, kpt.s)
ham.xc.apply_orbital_dependent_hamiltonian(kpt, psit_xG, Htpsit_xG,
ham.dH_asp)
self.timer.stop('Apply hamiltonian')
def get_pseudo_partial_waves(self):
phit_aj = [
setup.get_partial_waves_for_atomic_orbitals()
for setup in self.setups
]
return LFC(self.gd, phit_aj, kd=self.kd, cut=True, dtype=self.dtype)
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]
for f, psit_G in zip(f_n, kpt.psit_nG):
# Same as nt_G += f * abs(psit_G)**2, but much faster:
_gpaw.add_to_density(f, psit_G, 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 not isinstance(self.kpt_u[0].psit_nG, np.ndarray):
return None
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.kptband_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_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
def _get_wave_function_array(self, u, n, realspace=True, periodic=False):
assert realspace
kpt = self.kpt_u[u]
psit_G = kpt.psit_nG[n]
if periodic and self.dtype == complex:
k_c = self.kd.ibzk_kc[kpt.k]
return self.gd.plane_wave(-k_c) * psit_G
return psit_G
def write(self, writer, write_wave_functions=False):
FDPWWaveFunctions.write(self, writer)
if not write_wave_functions:
return
writer.add_array('values',
(self.nspins, self.kd.nibzkpts, self.bd.nbands) +
tuple(self.gd.get_size_of_global_array()), self.dtype)
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 * Bohr**-1.5)
def read(self, reader):
FDPWWaveFunctions.read(self, reader)
if 'values' not in reader.wave_functions:
return
c = reader.bohr**1.5
if reader.version < 0:
c = 1 # old gpw file
for kpt in self.kpt_u:
# 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:
psit_nG = reader.wave_functions.proxy('values', kpt.s, kpt.k)
psit_nG.scale = c
kpt.psit = UniformGridWaveFunctions(self.bd.nbands,
self.gd,
self.dtype,
psit_nG,
kpt=kpt.q,
dist=(self.bd.comm,
self.bd.comm.size),
spin=kpt.s,
collinear=True)
if self.world.size > 1:
# Read to memory:
for kpt in self.kpt_u:
kpt.psit.read_from_file()
def initialize_from_lcao_coefficients(self, basis_functions):
for kpt in self.mykpts:
kpt.psit = UniformGridWaveFunctions(self.bd.nbands,
self.gd,
self.dtype,
kpt=kpt.q,
dist=(self.bd.comm,
self.bd.comm.size, 1),
spin=kpt.s,
collinear=True)
kpt.psit_nG[:] = 0.0
mynbands = len(kpt.C_nM)
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)
class ExtraVacuumPoissonSolver(_PoissonSolver):
"""Wrapper around PoissonSolver extending the vacuum size.
Poisson equation is solved on the large grid defined by `gpts`
using `poissonsolver_large`.
If `coarses` is given, then the large grid is first coarsed
`coarses` times from to the original grid. Then, the coarse
potential is used to correct the boundary conditions
of the potential calculated on the original (small, fine)
grid by `poissonsolver_small`.
The parameters `nn_*` control the finite difference stencils
used in the coarsening, refining, and Laplace.
If the parameter `use_aux_grid` is `True`, an auxiliary
medium-sized grid is used between the large and small grids.
The parameter does not affect the result but can be used to
achieve speed-up depending on the grid sizes.
"""
def __init__(self, gpts, poissonsolver_large,
poissonsolver_small=None, coarses=0,
nn_coarse=3, nn_refine=3, nn_laplace=3,
use_aux_grid=True):
# TODO: Alternative options: vacuum size and h
self.N_large_fine_c = np.array(gpts, dtype=int)
self.Ncoar = coarses # coar == coarse
if self.Ncoar > 0:
self.use_coarse = True
else:
self.use_coarse = False
self.ps_large_coar = create_poisson_solver(poissonsolver_large)
if self.use_coarse:
self.ps_small_fine = create_poisson_solver(poissonsolver_small)
else:
assert poissonsolver_small is None
self.nn_coarse = nn_coarse
self.nn_refine = nn_refine
self.nn_laplace = nn_laplace
self.use_aux_grid = use_aux_grid
self._initialized = False
def set_grid_descriptor(self, gd):
# If non-periodic boundary conditions is used,
# there is problems with auxiliary grid.
# Maybe with use_aux_grid=False it would work?
if gd.pbc_c.any():
raise NotImplementedError('Only non-periodic boundary '
'conditions are tested')
self.gd_small_fine = gd
assert np.all(self.gd_small_fine.N_c <= self.N_large_fine_c), \
'extended grid has to be larger than the original one'
if self.use_coarse:
# 1.1. Construct coarse chain on the small grid
self.coarser_i = []
gd = self.gd_small_fine
N_c = self.N_large_fine_c.copy()
for i in range(self.Ncoar):
gd2 = gd.coarsen()
self.coarser_i.append(Transformer(gd, gd2, self.nn_coarse))
N_c //= 2
gd = gd2
self.gd_small_coar = gd
else:
self.gd_small_coar = self.gd_small_fine
N_c = self.N_large_fine_c
# 1.2. Construct coarse extended grid
self.gd_large_coar = ext_gd(self.gd_small_coar, N_c=N_c)
# Initialize poissonsolvers
self.ps_large_coar.set_grid_descriptor(self.gd_large_coar)
if not self.use_coarse:
return
self.ps_small_fine.set_grid_descriptor(self.gd_small_fine)
if self.use_aux_grid:
# 2.1. Construct an auxiliary grid that is the small grid plus
# a buffer region allowing Laplace and refining
# with the used stencils
buf = self.nn_refine
for i in range(self.Ncoar):
buf = 2 * buf + self.nn_refine
buf += self.nn_laplace
div = 2**self.Ncoar
if buf % div != 0:
buf += div - buf % div
N_c = self.gd_small_fine.N_c + 2 * buf
if np.any(N_c > self.N_large_fine_c):
self.use_aux_grid = False
N_c = self.N_large_fine_c
self.gd_aux_fine = ext_gd(self.gd_small_fine, N_c=N_c)
else:
self.gd_aux_fine = ext_gd(self.gd_small_fine,
N_c=self.N_large_fine_c)
# 2.2 Construct Laplace on the aux grid
self.laplace_aux_fine = Laplace(self.gd_aux_fine, - 0.25 / np.pi,
self.nn_laplace)
# 2.3 Construct refine chain
self.refiner_i = []
gd = self.gd_aux_fine
N_c = gd.N_c.copy()
for i in range(self.Ncoar):
gd2 = gd.coarsen()
self.refiner_i.append(Transformer(gd2, gd, self.nn_refine))
N_c //= 2
gd = gd2
self.refiner_i = self.refiner_i[::-1]
self.gd_aux_coar = gd
if self.use_aux_grid:
# 2.4 Construct large coarse grid from aux grid
self.gd_large_coar_from_aux = ext_gd(self.gd_aux_coar,
N_c=self.gd_large_coar.N_c)
# Check the consistency of the grids
gd1 = self.gd_large_coar
gd2 = self.gd_large_coar_from_aux
assert np.all(gd1.N_c == gd2.N_c) and np.all(gd1.h_cv == gd2.h_cv)
self._initialized = False
def _init(self):
if self._initialized:
return
# Allocate arrays
self.phi_large_coar_g = self.gd_large_coar.zeros()
self._initialized = True
# Initialize poissonsolvers
#self.ps_large_coar._init()
#if not self.use_coarse:
# return
#self.ps_small_fine._init()
def solve(self, phi, rho, **kwargs):
self._init()
phi_small_fine_g = phi
rho_small_fine_g = rho.copy()
if self.use_coarse:
# 1.1. Coarse rho on the small grid
tmp_g = rho_small_fine_g
for coarser in self.coarser_i:
tmp_g = coarser.apply(tmp_g)
rho_small_coar_g = tmp_g
else:
rho_small_coar_g = rho_small_fine_g
# 1.2. Extend rho to the large grid
rho_large_coar_g = self.gd_large_coar.zeros()
extend_array(self.gd_small_coar, self.gd_large_coar,
rho_small_coar_g, rho_large_coar_g)
# 1.3 Solve potential on the large coarse grid
niter_large = self.ps_large_coar.solve(self.phi_large_coar_g,
rho_large_coar_g, **kwargs)
rho_large_coar_g = None
if not self.use_coarse:
deextend_array(self.gd_small_fine, self.gd_large_coar,
phi_small_fine_g, self.phi_large_coar_g)
return niter_large
if self.use_aux_grid:
# 2.1 De-extend the potential to the auxiliary grid
phi_aux_coar_g = self.gd_aux_coar.empty()
deextend_array(self.gd_aux_coar, self.gd_large_coar_from_aux,
phi_aux_coar_g, self.phi_large_coar_g)
else:
phi_aux_coar_g = self.phi_large_coar_g
# 3.1 Refine the potential
tmp_g = phi_aux_coar_g
for refiner in self.refiner_i:
tmp_g = refiner.apply(tmp_g)
phi_aux_coar_g = None
phi_aux_fine_g = tmp_g
# 3.2 Calculate the corresponding density with Laplace
# (the refined coarse density would not accurately match with
# the potential)
rho_aux_fine_g = self.gd_aux_fine.empty()
self.laplace_aux_fine.apply(phi_aux_fine_g, rho_aux_fine_g)
# 3.3 De-extend the potential and density to the small grid
cor_phi_small_fine_g = self.gd_small_fine.empty()
deextend_array(self.gd_small_fine, self.gd_aux_fine,
cor_phi_small_fine_g, phi_aux_fine_g)
phi_aux_fine_g = None
cor_rho_small_fine_g = self.gd_small_fine.empty()
deextend_array(self.gd_small_fine, self.gd_aux_fine,
cor_rho_small_fine_g, rho_aux_fine_g)
rho_aux_fine_g = None
# 3.4 Remove the correcting density and potential
rho_small_fine_g -= cor_rho_small_fine_g
phi_small_fine_g -= cor_phi_small_fine_g
# 3.5 Solve potential on the small grid
niter_small = self.ps_small_fine.solve(phi_small_fine_g,
rho_small_fine_g, **kwargs)
# 3.6 Correct potential and density
phi_small_fine_g += cor_phi_small_fine_g
# rho_small_fine_g += cor_rho_small_fine_g
return (niter_large, niter_small)
def estimate_memory(self, mem):
self.ps_large_coar.estimate_memory(mem.subnode('Large grid Poisson'))
if self.use_coarse:
ps = self.ps_small_fine
ps.estimate_memory(mem.subnode('Small grid Poisson'))
mem.subnode('Large coarse phi', self.gd_large_coar.bytecount())
tmp = max(self.gd_small_fine.bytecount(),
self.gd_large_coar.bytecount())
if self.use_coarse:
tmp = max(tmp,
self.gd_aux_coar.bytecount(),
self.gd_aux_fine.bytecount() * 2 +
self.gd_small_fine.bytecount(),
self.gd_aux_fine.bytecount() +
self.gd_small_fine.bytecount() * 2)
mem.subnode('Temporary arrays', tmp)
def get_description(self):
line = '%s with ' % self.__class__.__name__
if self.use_coarse:
line += 'large and small grids'
else:
line += 'large grid'
lines = [line]
def add_line(line, pad=0):
lines.extend(['%s%s' % (' ' * pad, line)])
def get_cell(ps):
descr = ps.get_description().replace('\n', '\n%s' % (' ' * 8))
add_line('Poisson solver: %s' % descr, 8)
if hasattr(ps, 'gd'):
gd = ps.gd
par = cell_to_cellpar(gd.cell_cv * Bohr)
h_eff = gd.dv**(1.0 / 3.0) * Bohr
l1 = '{:8d} x {:8d} x {:8d} points'.format(*gd.N_c)
l2 = '{:8.2f} x {:8.2f} x {:8.2f} AA'.format(*par[:3])
l3 = 'Effective grid spacing dv^(1/3) = {:.4f}'.format(h_eff)
add_line('Grid: %s' % l1, 8)
add_line(' %s' % l2, 8)
add_line(' %s' % l3, 8)
add_line('Large grid:', 4)
get_cell(self.ps_large_coar)
if self.use_coarse:
add_line('Small grid:', 4)
get_cell(self.ps_small_fine)
return '\n'.join(lines)
def todict(self):
d = {'name': self.__class__.__name__}
d['gpts'] = self.N_large_fine_c
d['coarses'] = self.Ncoar
d['nn_coarse'] = self.nn_coarse
d['nn_refine'] = self.nn_refine
d['nn_laplace'] = self.nn_laplace
d['use_aux_grid'] = self.use_aux_grid
d['poissonsolver_large'] = self.ps_large_coar.todict()
if self.use_coarse:
d['poissonsolver_small'] = self.ps_small_fine.todict()
else:
d['poissonsolver_small'] = None
return d