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:
kpt.psit_nG = reader.wave_functions.proxy('values', kpt.s, kpt.k)
kpt.psit_nG.scale = c
if self.world.size == 1:
return
# Read to memory:
for kpt in self.kpt_u:
psit_nG = kpt.psit_nG
kpt.psit_nG = self.empty(self.bd.mynbands)
# Read band by band to save memory
for myn, psit_G in enumerate(kpt.psit_nG):
n = self.bd.global_index(myn)
# XXX number of bands could have been rounded up!
if n >= len(psit_nG):
break
if self.gd.comm.rank == 0:
big_psit_G = np.asarray(psit_nG[n], self.dtype)
else:
big_psit_G = None
self.gd.distribute(big_psit_G, psit_G)
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 __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_setups(self, setups):
self.pd = PWDescriptor(self.ecut, self.gd, self.kd.ibzk_qc)
pt = LFC(self.gd, [setup.pt_j for setup in setups],
self.kpt_comm, dtype=self.dtype, forces=True)
self.pt = PWLFC(pt, self.pd)
FDPWWaveFunctions.set_setups(self, setups)
def initialize_wave_functions_from_basis_functions(self, basis_functions,
density, hamiltonian,
spos_ac):
FDPWWaveFunctions.initialize_wave_functions_from_basis_functions(
self, basis_functions, density, hamiltonian, spos_ac)
for kpt in self.kpt_u:
kpt.psit_nG = self.pd.fft(kpt.psit_nG)
def initialize_wave_functions_from_basis_functions(self, basis_functions,
density, hamiltonian,
spos_ac):
FDPWWaveFunctions.initialize_wave_functions_from_basis_functions(
self, basis_functions, density, hamiltonian, spos_ac)
for kpt in self.kpt_u:
kpt.psit_nG = self.pd.fft(kpt.psit_nG)
def set_setups(self, setups):
self.pd = PWDescriptor(self.ecut, self.gd, self.kd.ibzk_qc)
pt = LFC(self.gd, [setup.pt_j for setup in setups],
self.kpt_comm,
dtype=self.dtype,
forces=True)
self.pt = PWLFC(pt, self.pd)
FDPWWaveFunctions.set_setups(self, setups)
def __init__(self, ecut, diagksl, orthoksl, initksl, gd, nvalence, setups,
bd, world, kd, timer):
self.ecut = ecut / units.Hartree
# Set dtype=complex and gamma=False:
kd.gamma = False
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, complex, world, kd,
timer)
orthoksl.gd = self.pd
self.matrixoperator = MatrixOperator(orthoksl)
self.wd = self.pd
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 __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 __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, ecut, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd,
world, kd, timer):
self.ecut = ecut / units.Hartree
# Set dtype=complex and gamma=False:
kd.gamma = False
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd, complex,
world, kd, timer)
orthoksl.gd = self.pd
self.matrixoperator = MatrixOperator(orthoksl)
self.wd = self.pd
def __init__(self, ecut, fftwflags,
diagksl, orthoksl, initksl,
gd, nvalence, setups, bd, dtype,
world, kd, kptband_comm, timer):
self.ecut = ecut
self.fftwflags = fftwflags
self.ng_k = None # number of G-vectors for all IBZ k-points
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd, dtype,
world, kd, kptband_comm, timer)
self.orthoksl.gd = self.pd
self.matrixoperator = MatrixOperator(self.orthoksl)
def __init__(self, ecut, fftwflags,
diagksl, orthoksl, initksl,
gd, nvalence, setups, bd, dtype,
world, kd, timer):
self.ecut = ecut
self.fftwflags = fftwflags
self.ng_k = None # number of G-vectors for all IBZ k-points
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl,
gd, nvalence, setups, bd, dtype,
world, kd, timer)
self.orthoksl.gd = self.pd
self.matrixoperator = MatrixOperator(self.orthoksl)
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 set_setups(self, setups):
self.timer.start('PWDescriptor')
self.pd = PWDescriptor(self.ecut, self.gd, self.dtype, self.kd,
self.fftwflags)
self.timer.stop('PWDescriptor')
# Build array of number of plane wave coefficiants for all k-points
# in the IBZ:
self.ng_k = np.zeros(self.kd.nibzkpts, dtype=int)
for kpt in self.kpt_u:
if kpt.s == 0:
self.ng_k[kpt.k] = len(self.pd.Q_qG[kpt.q])
self.kd.comm.sum(self.ng_k)
self.pt = PWLFC([setup.pt_j for setup in setups], self.pd)
FDPWWaveFunctions.set_setups(self, setups)
def get_wave_function_array(self, n, k, s, realspace=True,
cut=True):
psit_G = FDPWWaveFunctions.get_wave_function_array(self, n, k, s,
realspace)
if cut and psit_G is not None and not realspace:
psit_G = psit_G[:self.ng_k[k]].copy()
return psit_G
def get_wave_function_array(self, n, k, s, realspace=True,
cut=True):
psit_G = FDPWWaveFunctions.get_wave_function_array(self, n, k, s,
realspace)
if cut and psit_G is not None and not realspace:
psit_G = psit_G[:self.ng_k[k]].copy()
return psit_G
def set_setups(self, setups):
self.timer.start('PWDescriptor')
self.pd = PWDescriptor(self.ecut, self.gd, self.dtype, self.kd,
self.fftwflags)
self.timer.stop('PWDescriptor')
# Build array of number of plane wave coefficiants for all k-points
# in the IBZ:
self.ng_k = np.zeros(self.kd.nibzkpts)
for kpt in self.kpt_u:
if kpt.s == 0:
self.ng_k[kpt.k] = len(self.pd.Q_qG[kpt.q])
self.kd.comm.sum(self.ng_k)
self.pt = PWLFC([setup.pt_j for setup in setups], self.pd)
FDPWWaveFunctions.set_setups(self, setups)
def write(self, writer, write_wave_functions=False):
FDPWWaveFunctions.write(self, writer)
if not write_wave_functions:
return
writer.add_array(
'coefficients',
(self.nspins, self.kd.nibzkpts, self.bd.nbands, self.pd.ngmax),
complex)
c = units.Bohr**-1.5
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,
realspace=False,
cut=False)
writer.fill(psit_G * c)
writer.add_array('indices', (self.kd.nibzkpts, self.pd.ngmax),
np.int32)
if self.bd.comm.rank > 0:
return
Q_G = np.empty(self.pd.ngmax, np.int32)
kk = 0
for r in range(self.kd.comm.size):
for q, ks in enumerate(self.kd.get_indices(r)):
s, k = divmod(ks, self.kd.nibzkpts)
ng = self.ng_k[k]
if s == 1:
return
if r == self.kd.comm.rank:
Q_G[:ng] = self.pd.Q_qG[q]
if r > 0:
self.kd.comm.send(Q_G, 0)
if self.kd.comm.rank == 0:
if r > 0:
self.kd.comm.receive(Q_G, r)
Q_G[ng:] = -1
writer.fill(Q_G)
assert k == kk
kk += 1
def __str__(self):
s = 'Wave functions: Plane wave expansion\n'
s += ' Cutoff energy: %.3f eV\n' % (self.pd.ecut * units.Hartree)
if self.dtype == float:
s += (' Number of coefficients: %d (reduced to %d)\n' %
(self.pd.ngmax * 2 - 1, self.pd.ngmax))
else:
s += (' Number of coefficients (min, max): %d, %d\n' %
(self.pd.ngmin, self.pd.ngmax))
if fftw.FFTPlan is fftw.NumpyFFTPlan:
s += " Using Numpy's FFT\n"
else:
s += ' Using FFTW library\n'
return s + FDPWWaveFunctions.__str__(self)
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 read(self, reader):
FDPWWaveFunctions.read(self, reader)
if 'coefficients' not in reader.wave_functions:
return
Q_kG = reader.wave_functions.indices
for kpt in self.kpt_u:
if kpt.s == 0:
Q_G = Q_kG[kpt.k]
ng = self.ng_k[kpt.k]
assert (Q_G[:ng] == self.pd.Q_qG[kpt.q]).all()
assert (Q_G[ng:] == -1).all()
c = reader.bohr**1.5
if reader.version < 0:
c = 1 # old gpw file
for kpt in self.kpt_u:
ng = self.ng_k[kpt.k]
kpt.psit_nG = reader.wave_functions.proxy('coefficients',
kpt.s, kpt.k)
kpt.psit_nG.scale = c
kpt.psit_nG.length_of_last_dimension = ng
if self.world.size == 1:
return
# Read to memory:
for kpt in self.kpt_u:
psit_nG = kpt.psit_nG
kpt.psit_nG = self.empty(self.bd.mynbands, q=kpt.q)
ng = self.ng_k[kpt.k]
for myn, psit_G in enumerate(kpt.psit_nG):
n = self.bd.global_index(myn)
# XXX number of bands could have been rounded up!
if n < len(psit_nG):
psit_G[:] = psit_nG[n]
def _get_wave_function_array(self, u, n, realspace=True, phase=None):
psit_G = FDPWWaveFunctions._get_wave_function_array(self, u, n,
realspace)
if not realspace:
zeropadded_G = np.zeros(self.pd.ngmax, complex)
zeropadded_G[:len(psit_G)] = psit_G
return zeropadded_G
kpt = self.kpt_u[u]
if self.kd.gamma:
return self.pd.ifft(psit_G)
else:
if phase is None:
N_c = self.gd.N_c
k_c = self.kd.ibzk_kc[kpt.k]
eikr_R = np.exp(2j * pi * np.dot(np.indices(N_c).T,
k_c / N_c).T)
else:
eikr_R = phase
return self.pd.ifft(psit_G, kpt.q) * eikr_R
def _get_wave_function_array(self, u, n, realspace=True, phase=None):
psit_G = FDPWWaveFunctions._get_wave_function_array(self, u, n,
realspace)
if not realspace:
zeropadded_G = np.zeros(self.pd.ngmax, complex)
zeropadded_G[:len(psit_G)] = psit_G
return zeropadded_G
kpt = self.kpt_u[u]
if self.kd.gamma:
return self.pd.ifft(psit_G)
else:
if phase is None:
N_c = self.gd.N_c
k_c = self.kd.ibzk_kc[kpt.k]
eikr_R = np.exp(2j * pi * np.dot(np.indices(N_c).T,
k_c / N_c).T)
else:
eikr_R = phase
return self.pd.ifft(psit_G, kpt.q) * eikr_R
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
self.pd.estimate_memory(mem.subnode('PW-descriptor'))
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
self.kin.estimate_memory(mem.subnode('Kinetic operator'))
def set_positions(self, spos_ac):
if not self.kin.is_allocated():
self.kin.allocate()
FDPWWaveFunctions.set_positions(self, spos_ac)
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
self.kin.estimate_memory(mem.subnode('Kinetic operator'))
def estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
self.pd.estimate_memory(mem.subnode('PW-descriptor'))
def set_positions(self, spos_ac):
FDPWWaveFunctions.set_positions(self, spos_ac)
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 estimate_memory(self, mem):
FDPWWaveFunctions.estimate_memory(self, mem)
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 set_positions(self, spos_ac, atom_partition=None):
FDPWWaveFunctions.set_positions(self, spos_ac, atom_partition)
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 set_positions(self, spos_ac):
if not self.kin.is_allocated():
self.kin.allocate()
FDPWWaveFunctions.set_positions(self, spos_ac)
