1import numpy as np
2
3
4class DirectLCAO(object):
5    """Eigensolver for LCAO-basis calculation"""
6
7    def __init__(self, diagonalizer=None):
8        self.diagonalizer = diagonalizer
9        # ??? why should we be able to set
10        # this diagonalizer in both constructor and initialize?
11        self.has_initialized = False  # XXX
12
13    def initialize(self, gd, dtype, nao, diagonalizer=None):
14        self.gd = gd
15        self.nao = nao
16        if diagonalizer is not None:
17            self.diagonalizer = diagonalizer
18        assert self.diagonalizer is not None
19        self.has_initialized = True  # XXX
20
21    def reset(self):
22        pass
23
24    @property
25    def error(self):
26        return 0.0
27
28    @error.setter
29    def error(self, e):
30        pass
31
32    def calculate_hamiltonian_matrix(self, hamiltonian, wfs, kpt, Vt_xMM=None,
33                                     root=-1, add_kinetic=True):
34        # XXX document parallel stuff, particularly root parameter
35        assert self.has_initialized
36
37        bfs = wfs.basis_functions
38
39        # distributed_atomic_correction works with ScaLAPACK/BLACS in general.
40        # If SL is not enabled, it will not work with band parallelization.
41        # But no one would want that for a practical calculation anyway.
42        # dH_asp = wfs.atomic_correction.redistribute(wfs, hamiltonian.dH_asp)
43        # XXXXX fix atomic corrections
44        dH_asp = hamiltonian.dH_asp
45
46        if Vt_xMM is None:
47            wfs.timer.start('Potential matrix')
48            vt_G = hamiltonian.vt_sG[kpt.s]
49            Vt_xMM = bfs.calculate_potential_matrices(vt_G)
50            wfs.timer.stop('Potential matrix')
51
52        if bfs.gamma and wfs.dtype == float:
53            yy = 1.0
54            H_MM = Vt_xMM[0]
55        else:
56            wfs.timer.start('Sum over cells')
57            yy = 0.5
58            k_c = wfs.kd.ibzk_qc[kpt.q]
59            H_MM = (0.5 + 0.0j) * Vt_xMM[0]
60            for sdisp_c, Vt_MM in zip(bfs.sdisp_xc[1:], Vt_xMM[1:]):
61                H_MM += np.exp(2j * np.pi * np.dot(sdisp_c, k_c)) * Vt_MM
62            wfs.timer.stop('Sum over cells')
63
64        # Add atomic contribution
65        #
66        #           --   a     a  a*
67        # H      += >   P    dH  P
68        #  mu nu    --   mu i  ij nu j
69        #           aij
70        #
71        name = wfs.atomic_correction.__class__.__name__
72        wfs.timer.start(name)
73        wfs.atomic_correction.calculate_hamiltonian(kpt, dH_asp, H_MM, yy)
74        wfs.timer.stop(name)
75
76        wfs.timer.start('Distribute overlap matrix')
77        H_MM = wfs.ksl.distribute_overlap_matrix(
78            H_MM, root, add_hermitian_conjugate=(yy == 0.5))
79        wfs.timer.stop('Distribute overlap matrix')
80
81        if add_kinetic:
82            H_MM += wfs.T_qMM[kpt.q]
83        return H_MM
84
85    def iterate(self, hamiltonian, wfs, occ=None):
86        wfs.timer.start('LCAO eigensolver')
87
88        s = -1
89        for kpt in wfs.kpt_u:
90            if kpt.s != s:
91                s = kpt.s
92                wfs.timer.start('Potential matrix')
93                Vt_xMM = wfs.basis_functions.calculate_potential_matrices(
94                    hamiltonian.vt_sG[s])
95                wfs.timer.stop('Potential matrix')
96            self.iterate_one_k_point(hamiltonian, wfs, kpt, Vt_xMM)
97
98        wfs.timer.stop('LCAO eigensolver')
99
100    def iterate_one_k_point(self, hamiltonian, wfs, kpt, Vt_xMM):
101        if wfs.bd.comm.size > 1 and wfs.bd.strided:
102            raise NotImplementedError
103
104        H_MM = self.calculate_hamiltonian_matrix(hamiltonian, wfs, kpt, Vt_xMM,
105                                                 root=0)
106
107        # Decomposition step of overlap matrix can be skipped if we have
108        # cached the result and if the solver supports it (Elpa)
109        may_decompose = self.diagonalizer.accepts_decomposed_overlap_matrix
110        if may_decompose:
111            S_MM = wfs.decomposed_S_qMM[kpt.q]
112            is_already_decomposed = (S_MM is not None)
113            if S_MM is None:
114                # Contents of S_MM will be overwritten by eigensolver below.
115                S_MM = wfs.decomposed_S_qMM[kpt.q] = wfs.S_qMM[kpt.q].copy()
116        else:
117            is_already_decomposed = False
118            S_MM = wfs.S_qMM[kpt.q]
119
120        if kpt.eps_n is None:
121            kpt.eps_n = np.empty(wfs.bd.mynbands)
122
123        diagonalization_string = repr(self.diagonalizer)
124        wfs.timer.start(diagonalization_string)
125        # May overwrite S_MM (so the results will be stored as decomposed)
126        if kpt.C_nM is None:
127            kpt.C_nM = wfs.bd.empty(wfs.setups.nao, dtype=wfs.dtype)
128
129        self.diagonalizer.diagonalize(H_MM, kpt.C_nM, kpt.eps_n, S_MM,
130                                      is_already_decomposed)
131        wfs.timer.stop(diagonalization_string)
132
133        wfs.timer.start('Calculate projections')
134        # P_ani are not strictly necessary as required quantities can be
135        # evaluated directly using P_aMi/Paaqim.  We should perhaps get rid
136        # of the places in the LCAO code using P_ani directly
137        wfs.atomic_correction.calculate_projections(wfs, kpt)
138        wfs.timer.stop('Calculate projections')
139
140    def __repr__(self):
141        # The "diagonalizer" must be equal to the Kohn-Sham layout
142        # object presently.  That information will be printed in the
143        # text output anyway so we do not need it here.
144        #
145        # Although maybe it may be better to print it anyway...
146        return 'LCAO using direct dense diagonalizer'
147
148    def estimate_memory(self, mem, dtype):
149        pass
150        # self.diagonalizer.estimate_memory(mem, dtype) #XXX enable this
151