def __init__(self, ksl, gd, nvalence, setups, bd,
dtype, world, kd, kptband_comm, timer,
atomic_correction=None, collinear=True):
WaveFunctions.__init__(self, gd, nvalence, setups, bd,
dtype, collinear, world, kd,
kptband_comm, timer)
self.ksl = ksl
self.S_qMM = None
self.T_qMM = None
self.P_aqMi = None
self.debug_tci = False
if atomic_correction is None:
if ksl.using_blacs:
atomic_correction = 'scipy'
else:
atomic_correction = 'dense'
if isinstance(atomic_correction, str):
atomic_correction = get_atomic_correction(atomic_correction)
self.atomic_correction = atomic_correction
#self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma)
with self.timer('TCI: Evaluate splines'):
self.tciexpansions = TCIExpansions.new_from_setups(setups)
self.basis_functions = BasisFunctions(gd,
[setup.phit_j
for setup in setups],
kd,
dtype=dtype,
cut=True)
def __init__(self,
ksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
kptband_comm,
timer,
atomic_hamiltonian=None):
WaveFunctions.__init__(self, gd, nvalence, setups, bd, dtype, world,
kd, kptband_comm, timer)
self.ksl = ksl
self.S_qMM = None
self.T_qMM = None
self.P_aqMi = None
if atomic_hamiltonian is None:
if ksl.using_blacs:
atomic_hamiltonian = 'distributed'
else:
atomic_hamiltonian = 'dense'
if isinstance(atomic_hamiltonian, str):
atomic_hamiltonian = get_atomic_hamiltonian(atomic_hamiltonian)
self.atomic_hamiltonian = atomic_hamiltonian
self.timer.start('TCI: Evaluate splines')
self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma)
self.timer.stop('TCI: Evaluate splines')
self.basis_functions = BasisFunctions(
gd, [setup.phit_j for setup in setups], kd, dtype=dtype, cut=True)
def __init__(self, diagksl, orthoksl, initksl, *args, **kwargs):
WaveFunctions.__init__(self, *args, **kwargs)
self.diagksl = diagksl
self.orthoksl = orthoksl
self.initksl = initksl
self.set_orthonormalized(False)
self.overlap = Overlap(orthoksl, self.timer)
def __init__(self, diagksl, orthoksl, initksl, *args, **kwargs):
WaveFunctions.__init__(self, *args, **kwargs)
self.diagksl = diagksl
self.orthoksl = orthoksl
self.initksl = initksl
self.set_orthonormalized(False)
self.overlap = Overlap(orthoksl, self.timer)
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
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)
self.diagksl = None
self.orthoksl = BandLayouts(gd, bd, dtype)
self.initksl = None
self.overlap = None
self.rank_a = None
def calculate_atomic_density_matrices_with_occupation(self, D_asp, f_un):
ac = self.atomic_correction
if ac.implements_distributed_projections():
D2_asp = ac.redistribute(self, D_asp, type='asp', op='forth')
WaveFunctions.calculate_atomic_density_matrices_with_occupation(
self, D2_asp, f_un)
D3_asp = ac.redistribute(self, D2_asp, type='asp', op='back')
for a in D_asp:
D_asp[a][:] = D3_asp[a]
else:
WaveFunctions.calculate_atomic_density_matrices_with_occupation(
self, D_asp, f_un)
def read(self, reader):
WaveFunctions.read(self, reader)
if 'coefficients' not in reader.wave_functions:
return
for kpt in self.kpt_u:
C_nM = reader.wave_functions.proxy('coefficients', kpt.s, kpt.k)
kpt.C_nM = self.bd.empty(self.setups.nao, dtype=self.dtype)
for myn, C_M in enumerate(kpt.C_nM):
n = self.bd.global_index(myn)
# XXX number of bands could have been rounded up!
if n >= len(C_nM):
break
C_M[:] = C_nM[n] * Bohr**1.5
def write(self, writer, write_wave_functions=False):
WaveFunctions.write(self, writer)
if not write_wave_functions:
return
writer.add_array(
'coefficients',
(self.nspins, self.kd.nibzkpts, self.bd.nbands, self.setups.nao),
dtype=self.dtype)
for s in range(self.nspins):
for k in range(self.kd.nibzkpts):
C_nM = self.collect_array('C_nM', k, s)
writer.fill(C_nM * Bohr**-1.5)
def __init__(self, ksl, gd, nvalence, setups, bd,
dtype, world, kd, timer=None):
WaveFunctions.__init__(self, gd, nvalence, setups, bd,
dtype, world, kd, timer)
self.ksl = ksl
self.S_qMM = None
self.T_qMM = None
self.P_aqMi = None
self.timer.start('TCI: Evaluate splines')
self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma)
self.timer.stop('TCI: Evaluate splines')
self.basis_functions = BasisFunctions(gd,
[setup.phit_j
for setup in setups],
kd,
cut=True)
def set_positions(self, spos_ac, atom_partition=None):
move_wfs = (self.kpt_u[0].psit_nG is not None
and self.spos_ac is not None)
if move_wfs:
paste_wfs = self.wfs_mover.cut_wfs(self, spos_ac)
# This will update the positions -- and transfer, if necessary --
# the projection matrices which may be necessary for updating
# the wavefunctions.
WaveFunctions.set_positions(self, spos_ac, atom_partition)
if move_wfs and paste_wfs is not None:
paste_wfs()
self.set_orthonormalized(False)
self.pt.set_positions(spos_ac, atom_partition)
self.allocate_arrays_for_projections(self.pt.my_atom_indices)
self.positions_set = True
def __init__(self, ksl, gd, nvalence, setups, bd,
dtype, world, kd, timer=None):
WaveFunctions.__init__(self, gd, nvalence, setups, bd,
dtype, world, kd, timer)
self.ksl = ksl
self.S_qMM = None
self.T_qMM = None
self.P_aqMi = None
self.timer.start('TCI: Evaluate splines')
self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma)
self.timer.stop('TCI: Evaluate splines')
self.basis_functions = BasisFunctions(gd,
[setup.phit_j
for setup in setups],
kd.comm,
cut=True)
if not kd.gamma:
self.basis_functions.set_k_points(kd.ibzk_qc)
def collect_projections(self, P_ani):
if self.gd.comm.size == 1 and self.bd.comm.size == 1:
return np.concatenate([P_ni.T for P_ni in P_ani.values()])
assert len(self.kpt_u) == 1
self.kpt_u[0].P_ani = P_ani
all_P_ni = WaveFunctions.collect_projections(self, 0, 0)
if self.world.rank == 0:
P_In = all_P_ni.T.copy()
else:
nproj = sum([setup.ni for setup in self.setups])
P_In = np.empty((nproj, self.bd.nbands), self.pt.dtype)
self.world.broadcast(P_In, 0)
return P_In
def collect_projections(self, P_ani):
if self.gd.comm.size == 1 and self.bd.comm.size == 1:
return np.concatenate([P_ni.T for P_ni in P_ani.values()])
assert len(self.kpt_u) == 1
self.kpt_u[0].P_ani = P_ani
all_P_ni = WaveFunctions.collect_projections(self, 0, 0)
if self.world.rank == 0:
P_In = all_P_ni.T.copy()
else:
nproj = sum([setup.ni for setup in self.setups])
P_In = np.empty((nproj, self.bd.nbands), self.pt.dtype)
self.world.broadcast(P_In, 0)
return P_In
def __init__(self, parallel, initksl, reuse_wfs_method=None, **kwargs):
WaveFunctions.__init__(self, **kwargs)
self.scalapack_parameters = parallel
self.initksl = initksl
if reuse_wfs_method is None or reuse_wfs_method == 'keep':
wfs_mover = NullWfsMover()
elif hasattr(reuse_wfs_method, 'cut_wfs'):
wfs_mover = reuse_wfs_method
elif reuse_wfs_method == 'paw':
wfs_mover = PseudoPartialWaveWfsMover()
elif reuse_wfs_method == 'lcao':
wfs_mover = LCAOWfsMover()
else:
raise ValueError(
'Strange way to reuse wfs: {}'.format(reuse_wfs_method))
self.wfs_mover = wfs_mover
self.set_orthonormalized(False)
self._work_matrix_nn = None # storage for H, S, ...
self._work_array = None
def set_positions(self, spos_ac):
self.timer.start('Basic WFS set positions')
WaveFunctions.set_positions(self, spos_ac)
self.timer.stop('Basic WFS set positions')
self.timer.start('Basis functions set positions')
self.basis_functions.set_positions(spos_ac)
self.timer.stop('Basis functions set positions')
if self.ksl is not None:
self.basis_functions.set_matrix_distribution(self.ksl.Mstart,
self.ksl.Mstop)
nq = len(self.kd.ibzk_qc)
nao = self.setups.nao
mynbands = self.mynbands
Mstop = self.ksl.Mstop
Mstart = self.ksl.Mstart
mynao = Mstop - Mstart
if self.ksl.using_blacs: # XXX
# S and T have been distributed to a layout with blacs, so
# discard them to force reallocation from scratch.
#
# TODO: evaluate S and T when they *are* distributed, thus saving
# memory and avoiding this problem
self.S_qMM = None
self.T_qMM = None
S_qMM = self.S_qMM
T_qMM = self.T_qMM
if S_qMM is None: # XXX
# First time:
assert T_qMM is None
if self.ksl.using_blacs: # XXX
self.tci.set_matrix_distribution(Mstart, mynao)
S_qMM = np.empty((nq, mynao, nao), self.dtype)
T_qMM = np.empty((nq, mynao, nao), self.dtype)
for kpt in self.kpt_u:
if kpt.C_nM is None:
kpt.C_nM = np.empty((mynbands, nao), self.dtype)
self.allocate_arrays_for_projections(
self.basis_functions.my_atom_indices)
self.P_aqMi = {}
for a in self.basis_functions.my_atom_indices:
ni = self.setups[a].ni
self.P_aqMi[a] = np.empty((nq, nao, ni), self.dtype)
for kpt in self.kpt_u:
q = kpt.q
kpt.P_aMi = dict([(a, P_qMi[q])
for a, P_qMi in self.P_aqMi.items()])
self.timer.start('TCI: Calculate S, T, P')
# Calculate lower triangle of S and T matrices:
self.tci.calculate(spos_ac, S_qMM, T_qMM, self.P_aqMi)
add_paw_correction_to_overlap(self.setups, self.P_aqMi, S_qMM,
self.ksl.Mstart, self.ksl.Mstop)
self.timer.stop('TCI: Calculate S, T, P')
S_MM = None # allow garbage collection of old S_qMM after redist
S_qMM = self.ksl.distribute_overlap_matrix(S_qMM)
T_qMM = self.ksl.distribute_overlap_matrix(T_qMM)
for kpt in self.kpt_u:
q = kpt.q
kpt.S_MM = S_qMM[q]
kpt.T_MM = T_qMM[q]
if (debug and self.band_comm.size == 1 and self.gd.comm.rank == 0 and
nao > 0 and not self.ksl.using_blacs):
# S and T are summed only on comm master, so check only there
from numpy.linalg import eigvalsh
self.timer.start('Check positive definiteness')
for S_MM in S_qMM:
tri2full(S_MM, UL='L')
smin = eigvalsh(S_MM).real.min()
if smin < 0:
raise RuntimeError('Overlap matrix has negative '
'eigenvalue: %e' % smin)
self.timer.stop('Check positive definiteness')
self.positions_set = True
self.S_qMM = S_qMM
self.T_qMM = T_qMM
def set_positions(self, spos_ac, atom_partition=None, move_wfs=False):
oldspos_ac = self.spos_ac
with self.timer('Basic WFS set positions'):
WaveFunctions.set_positions(self, spos_ac, atom_partition)
with self.timer('Basis functions set positions'):
self.basis_functions.set_positions(spos_ac)
if self.ksl is not None:
self.basis_functions.set_matrix_distribution(self.ksl.Mstart,
self.ksl.Mstop)
nq = len(self.kd.ibzk_qc)
nao = self.setups.nao
mynbands = self.bd.mynbands
Mstop = self.ksl.Mstop
Mstart = self.ksl.Mstart
mynao = Mstop - Mstart
#if self.ksl.using_blacs: # XXX
# S and T have been distributed to a layout with blacs, so
# discard them to force reallocation from scratch.
#
# TODO: evaluate S and T when they *are* distributed, thus saving
# memory and avoiding this problem
for kpt in self.kpt_u:
kpt.S_MM = None
kpt.T_MM = None
# Free memory in case of old matrices:
self.S_qMM = self.T_qMM = self.P_aqMi = None
if self.dtype == complex and oldspos_ac is not None:
update_phases([kpt.C_nM for kpt in self.kpt_u],
[kpt.q for kpt in self.kpt_u],
self.kd.ibzk_qc, spos_ac, oldspos_ac,
self.setups, Mstart)
for kpt in self.kpt_u:
if kpt.C_nM is None:
kpt.C_nM = np.empty((mynbands, nao), self.dtype)
if 0:#self.debug_tci:
#if self.ksl.using_blacs:
# self.tci.set_matrix_distribution(Mstart, mynao)
oldS_qMM = np.empty((nq, mynao, nao), self.dtype)
oldT_qMM = np.empty((nq, mynao, nao), self.dtype)
oldP_aqMi = {}
for a in self.basis_functions.my_atom_indices:
ni = self.setups[a].ni
oldP_aqMi[a] = np.empty((nq, nao, ni), self.dtype)
# Calculate lower triangle of S and T matrices:
self.timer.start('tci calculate')
#self.tci.calculate(spos_ac, oldS_qMM, oldT_qMM,
# oldP_aqMi)
self.timer.stop('tci calculate')
self.timer.start('mktci')
manytci = self.tciexpansions.get_manytci_calculator(
self.setups, self.gd, spos_ac, self.kd.ibzk_qc, self.dtype,
self.timer)
self.timer.stop('mktci')
self.manytci = manytci
self.newtci = manytci.tci
my_atom_indices = self.basis_functions.my_atom_indices
self.timer.start('ST tci')
newS_qMM, newT_qMM = manytci.O_qMM_T_qMM(self.gd.comm,
Mstart, Mstop,
self.ksl.using_blacs)
self.timer.stop('ST tci')
self.timer.start('P tci')
P_qIM = manytci.P_qIM(my_atom_indices)
self.timer.stop('P tci')
self.P_aqMi = newP_aqMi = manytci.P_aqMi(my_atom_indices)
self.P_qIM = P_qIM # XXX atomic correction
# TODO
# OK complex/conj, periodic images
# OK scalapack
# derivatives/forces
# sparse
# use symmetry/conj tricks to reduce calculations
# enable caching of spherical harmonics
#if self.atomic_correction.name != 'dense':
#from gpaw.lcao.newoverlap import newoverlap
#self.P_neighbors_a, self.P_aaqim = newoverlap(self, spos_ac)
self.atomic_correction.gobble_data(self)
#if self.atomic_correction.name == 'scipy':
# Pold_qIM = self.atomic_correction.Psparse_qIM
# for q in range(nq):
# maxerr = abs(Pold_qIM[q] - P_qIM[q]).max()
# print('sparse maxerr', maxerr)
# assert maxerr == 0
self.atomic_correction.add_overlap_correction(self, newS_qMM)
if self.debug_tci:
self.atomic_correction.add_overlap_correction(self, oldS_qMM)
if self.atomic_correction.implements_distributed_projections():
my_atom_indices = self.atomic_correction.get_a_values()
else:
my_atom_indices = self.basis_functions.my_atom_indices
self.allocate_arrays_for_projections(my_atom_indices)
#S_MM = None # allow garbage collection of old S_qMM after redist
if self.debug_tci:
oldS_qMM = self.ksl.distribute_overlap_matrix(oldS_qMM, root=-1)
oldT_qMM = self.ksl.distribute_overlap_matrix(oldT_qMM, root=-1)
newS_qMM = self.ksl.distribute_overlap_matrix(newS_qMM, root=-1)
newT_qMM = self.ksl.distribute_overlap_matrix(newT_qMM, root=-1)
#if (debug and self.bd.comm.size == 1 and self.gd.comm.rank == 0 and
# nao > 0 and not self.ksl.using_blacs):
# S and T are summed only on comm master, so check only there
# from numpy.linalg import eigvalsh
# self.timer.start('Check positive definiteness')
# for S_MM in S_qMM:
# tri2full(S_MM, UL='L')
# smin = eigvalsh(S_MM).real.min()
# if smin < 0:
# raise RuntimeError('Overlap matrix has negative '
# 'eigenvalue: %e' % smin)
# self.timer.stop('Check positive definiteness')
self.positions_set = True
if self.debug_tci:
Serr = np.abs(newS_qMM - oldS_qMM).max()
Terr = np.abs(newT_qMM - oldT_qMM).max()
print('S maxerr', Serr)
print('T maxerr', Terr)
try:
assert Terr < 1e-15, Terr
except AssertionError:
np.set_printoptions(precision=6)
if self.world.rank == 0:
print(newT_qMM)
print(oldT_qMM)
print(newT_qMM - oldT_qMM)
raise
assert Serr < 1e-15, Serr
assert len(oldP_aqMi) == len(newP_aqMi)
for a in oldP_aqMi:
Perr = np.abs(oldP_aqMi[a] - newP_aqMi[a]).max()
assert Perr < 1e-15, (a, Perr)
for kpt in self.kpt_u:
q = kpt.q
kpt.S_MM = newS_qMM[q]
kpt.T_MM = newT_qMM[q]
self.S_qMM = newS_qMM
self.T_qMM = newT_qMM
# Elpa wants to reuse the decomposed form of S_qMM.
# We need to keep track of the existence of that object here,
# since this is where we change S_qMM. Hence, expect this to
# become arrays after the first diagonalization:
self.decomposed_S_qMM = [None] * len(self.S_qMM)
def set_setups(self, setups):
WaveFunctions.set_setups(self, setups)
def set_setups(self, setups):
WaveFunctions.set_setups(self, setups)
def __str__(self):
s = 'Wave functions: LCAO\n'
s += ' Diagonalizer: %s\n' % self.ksl.get_description()
s += ' Atomic Correction: %s\n' % self.atomic_correction.description
s += ' Datatype: %s\n' % self.dtype.__name__
return s + WaveFunctions.__str__(self)
def set_eigensolver(self, eigensolver):
WaveFunctions.set_eigensolver(self, eigensolver)
eigensolver.initialize(self.gd, self.dtype, self.setups.nao, self.ksl)
def set_positions(self, spos_ac):
self.timer.start('Basic WFS set positions')
WaveFunctions.set_positions(self, spos_ac)
self.timer.stop('Basic WFS set positions')
self.timer.start('Basis functions set positions')
self.basis_functions.set_positions(spos_ac)
self.timer.stop('Basis functions set positions')
if self.ksl is not None:
self.basis_functions.set_matrix_distribution(self.ksl.Mstart,
self.ksl.Mstop)
nq = len(self.kd.ibzk_qc)
nao = self.setups.nao
mynbands = self.bd.mynbands
Mstop = self.ksl.Mstop
Mstart = self.ksl.Mstart
mynao = Mstop - Mstart
if self.ksl.using_blacs: # XXX
# S and T have been distributed to a layout with blacs, so
# discard them to force reallocation from scratch.
#
# TODO: evaluate S and T when they *are* distributed, thus saving
# memory and avoiding this problem
self.S_qMM = None
self.T_qMM = None
S_qMM = self.S_qMM
T_qMM = self.T_qMM
if S_qMM is None: # XXX
# First time:
assert T_qMM is None
if self.ksl.using_blacs: # XXX
self.tci.set_matrix_distribution(Mstart, mynao)
S_qMM = np.empty((nq, mynao, nao), self.dtype)
T_qMM = np.empty((nq, mynao, nao), self.dtype)
for kpt in self.kpt_u:
if kpt.C_nM is None:
kpt.C_nM = np.empty((mynbands, nao), self.dtype)
self.allocate_arrays_for_projections(
self.basis_functions.my_atom_indices)
self.P_aqMi = {}
for a in self.basis_functions.my_atom_indices:
ni = self.setups[a].ni
self.P_aqMi[a] = np.empty((nq, nao, ni), self.dtype)
for kpt in self.kpt_u:
q = kpt.q
kpt.P_aMi = dict([(a, P_qMi[q])
for a, P_qMi in self.P_aqMi.items()])
self.timer.start('TCI: Calculate S, T, P')
# Calculate lower triangle of S and T matrices:
self.tci.calculate(spos_ac, S_qMM, T_qMM, self.P_aqMi)
add_paw_correction_to_overlap(self.setups, self.P_aqMi, S_qMM,
self.ksl.Mstart, self.ksl.Mstop)
self.timer.stop('TCI: Calculate S, T, P')
S_MM = None # allow garbage collection of old S_qMM after redist
S_qMM = self.ksl.distribute_overlap_matrix(S_qMM)
T_qMM = self.ksl.distribute_overlap_matrix(T_qMM)
for kpt in self.kpt_u:
q = kpt.q
kpt.S_MM = S_qMM[q]
kpt.T_MM = T_qMM[q]
if (debug and self.band_comm.size == 1 and self.gd.comm.rank == 0 and
nao > 0 and not self.ksl.using_blacs):
# S and T are summed only on comm master, so check only there
from numpy.linalg import eigvalsh
self.timer.start('Check positive definiteness')
for S_MM in S_qMM:
tri2full(S_MM, UL='L')
smin = eigvalsh(S_MM).real.min()
if smin < 0:
raise RuntimeError('Overlap matrix has negative '
'eigenvalue: %e' % smin)
self.timer.stop('Check positive definiteness')
self.positions_set = True
self.S_qMM = S_qMM
self.T_qMM = T_qMM
def write(self, writer, write_wave_functions=False):
WaveFunctions.write(self, writer)
if write_wave_functions:
self.write_wave_functions(writer)
def read(self, reader):
WaveFunctions.read(self, reader)
r = reader.wave_functions
if 'coefficients' in r:
self.read_wave_functions(r)
def set_setups(self, setups):
WaveFunctions.set_setups(self, setups)
if not self.gamma:
self.pt.set_k_points(self.kd.ibzk_qc)
def set_setups(self, setups):
WaveFunctions.set_setups(self, setups)
if not self.gamma:
self.pt.set_k_points(self.kd.ibzk_qc)
def set_positions(self, spos_ac):
WaveFunctions.set_positions(self, spos_ac)
self.set_orthonormalized(False)
self.pt.set_positions(spos_ac)
self.allocate_arrays_for_projections(self.pt.my_atom_indices)
self.positions_set = True
def set_positions(self, spos_ac):
WaveFunctions.set_positions(self, spos_ac)
self.set_orthonormalized(False)
self.pt.set_positions(spos_ac)
self.allocate_arrays_for_projections(self.pt.my_atom_indices)
self.positions_set = True
def set_eigensolver(self, eigensolver):
WaveFunctions.set_eigensolver(self, eigensolver)
eigensolver.initialize(self.gd, self.dtype, self.setups.nao, self.ksl)
