1# coding: utf-8
2"""
3Flows for electron-phonon calculations (high-level interface)
4"""
5import numpy as np
6
7from abipy.core.kpoints import kpath_from_bounds_and_ndivsm
8from .works import Work, PhononWork, PhononWfkqWork
9from .flows import Flow
10
11
12class EphPotFlow(Flow):
13    r"""
14    This flow computes the e-ph scattering potentials on a q-mesh defined by ngqpt
15    and a list of q-points (usually a q-path) specified by the user.
16    The DFPT potentials on the q-mesh are merged in the DVDB located in the outdata
17    of the second work while the DFPT potentials on the q-path are merged in the DVDB
18    located in the outdata of the third work.
19    These DVDB files are then passed to the EPH code to compute the average over the unit
20    cell of the periodic part of the scattering potentials as a function of q.
21    Results are stored in the V1QAVG.nc files of the outdata of the tasks in the fourth work.
22    """
23
24    @classmethod
25    def from_scf_input(cls, workdir, scf_input, ngqpt, qbounds,
26                       ndivsm=5, with_becs=True, ddk_tolerance=None, prepgkk=0, manager=None):
27        """
28        Build the flow from an input file representing a GS calculation.
29
30        Args:
31            workdir: Working directory.
32            scf_input: Input for the GS SCF run.
33            ngqpt: 3 integers defining the q-mesh.
34            qbounds: List of boundaries defining the q-path used for the computation of the GKQ files.
35                The q-path is automatically generated using `ndivsm` and the reciprocal-space metric.
36                If `ndivsm` is 0, the code assumes that `qbounds` contains the full list of q-points
37                and no pre-processing is performed.
38            ndivsm: Number of points in the smallest segment of the path defined by `qbounds`.
39                Use 0 to pass list of q-points.
40            with_becs: Activate calculation of Electric field and Born effective charges.
41            ddk_tolerance: dict {"varname": value} with the tolerance used in the DDK run if `with_becs`.
42            prepgkk: 1 to activate computation of all 3 * natom perts (debugging option).
43            manager: |TaskManager| object.
44        """
45        flow = cls(workdir=workdir, manager=manager)
46
47        # First work with GS run.
48        scf_task = flow.register_scf_task(scf_input)[0]
49
50        # Second work to compute phonons on the input nqgpt q-mesh.
51        work_qmesh = PhononWork.from_scf_task(scf_task, qpoints=ngqpt, is_ngqpt=True,
52                                              with_becs=with_becs, ddk_tolerance=ddk_tolerance)
53        flow.register_work(work_qmesh)
54
55        if ndivsm > 0:
56            # Generate list of q-points from qbounds and ndivsm.
57            qpath_list = kpath_from_bounds_and_ndivsm(qbounds, ndivsm, scf_input.structure)
58        elif ndivsm == 0:
59            # Use input list of q-points.
60            qpath_list = np.reshape(qbounds, (-1, 3))
61        else:
62            raise ValueError("ndivsm cannot be negative. Received ndivsm: %s" % ndivsm)
63
64        # Third Work: compute WFK/WFQ and phonons for qpt in qpath_list.
65        # Don't include BECS because they have been already computed in the previous work.
66        work_qpath = PhononWfkqWork.from_scf_task(
67                       scf_task, qpath_list, ph_tolerance=None, tolwfr=1.0e-22, nband=None,
68                       with_becs=False, ddk_tolerance=None, shiftq=(0, 0, 0), is_ngqpt=False, remove_wfkq=True,
69                       prepgkk=prepgkk, manager=manager)
70
71        flow.register_work(work_qpath)
72
73        # Now we compute matrix elements fully ab-initio for each q-point.
74        eph_work = Work()
75
76        for eph_task in (-15, 15):
77            eph_inp = scf_input.new_with_vars(
78                optdriver=7,
79                ddb_ngqpt=ngqpt,    # q-mesh associated to the DDB file.
80                #dvdb_ngqpt=ngqpt,  # q-mesh associated to the DDVDB file.
81                prtphdos=0,
82                eph_task=eph_task
83            )
84
85            if eph_task == -15:
86                # Use DVDB with ab-initio POTS along q-path to produce V1QAVG
87                deps = {work_qmesh: "DDB", work_qpath: "DVDB"}
88            elif eph_task == 15:
89                # Use q-mesh to interpolate along the same q-path as above.
90                deps = {work_qmesh: ["DDB", "DVDB"]}
91                eph_inp.set_vars(ph_nqpath=len(qpath_list), ph_qpath=qpath_list)
92
93            eph_work.register_eph_task(eph_inp, deps=deps)
94
95        flow.register_work(eph_work)
96
97        return flow
98
99
100class GkqPathFlow(Flow):
101    r"""
102    This flow computes the gkq e-ph matrix elements <k+q|\Delta V_q|k> for a list of q-points (usually a q-path).
103    The results are stored in the GKQ.nc files for the different q-points. These files can be used to analyze the behaviour
104    of the e-ph matrix elements as a function of qpts with the the objects provided by the abipy.eph.gkq module.
105    It is also possible to compute the e-ph matrix elements using the interpolated DFPT potentials
106    if test_ft_interpolation is set to True.
107    """
108
109    @classmethod
110    def from_scf_input(cls, workdir, scf_input, ngqpt, qbounds,
111                       ndivsm=5, with_becs=True, ddk_tolerance=None,
112                       test_ft_interpolation=False, prepgkk=0, manager=None):
113        """
114        Build the flow from an input file representing a GS calculation.
115
116        Args:
117            workdir: Working directory.
118            scf_input: Input for the GS SCF run.
119            ngqpt: 3 integers defining the q-mesh.
120            qbounds: List of boundaries defining the q-path used for the computation of the GKQ files.
121                The q-path is automatically generated using `ndivsm` and the reciprocal-space metric.
122                If `ndivsm` is 0, the code assumes that `qbounds` contains the full list of q-points
123                and no pre-processing is performed.
124            ndivsm: Number of points in the smallest segment of the path defined by `qbounds`.
125                Use 0 to pass list of q-points.
126            with_becs: Activate calculation of Electric field and Born effective charges.
127            ddk_tolerance: dict {"varname": value} with the tolerance used in the DDK run if `with_becs`.
128            test_ft_interpolation: True to add an extra Work in which the GKQ files are computed
129                using the interpolated DFPT potentials and the q-mesh defined by `ngqpt`.
130                The quality of the interpolation depends on the convergence of the BECS, epsinf and `ngqpt`.
131            prepgkk: 1 to activate computation of all 3 * natom perts (debugging option).
132            manager: |TaskManager| object.
133        """
134        flow = cls(workdir=workdir, manager=manager)
135
136        # First work with GS run.
137        scf_task = flow.register_scf_task(scf_input)[0]
138
139        # Second work to compute phonons on the input nqgpt q-mesh.
140        work_qmesh = PhononWork.from_scf_task(scf_task, qpoints=ngqpt, is_ngqpt=True,
141                                              with_becs=with_becs, ddk_tolerance=ddk_tolerance)
142        flow.register_work(work_qmesh)
143
144        if ndivsm > 0:
145            # Generate list of q-points from qbounds and ndivsm.
146            qpath_list = kpath_from_bounds_and_ndivsm(qbounds, ndivsm, scf_input.structure)
147        elif ndivsm == 0:
148            # Use input list of q-points.
149            qpath_list = np.reshape(qbounds, (-1, 3))
150        else:
151            raise ValueError("ndivsm cannot be negative. Received ndivsm: %s" % ndivsm)
152
153        # Third Work. Compute WFK/WFQ and phonons for qpt in qpath_list.
154        # Don't include BECS because they have been already computed in the previous work.
155        work_qpath = PhononWfkqWork.from_scf_task(
156                       scf_task, qpath_list, ph_tolerance=None, tolwfr=1.0e-22, nband=None,
157                       with_becs=False, ddk_tolerance=None, shiftq=(0, 0, 0), is_ngqpt=False, remove_wfkq=False,
158                       prepgkk=prepgkk, manager=manager)
159
160        flow.register_work(work_qpath)
161
162        def make_eph_input(scf_inp, ngqpt, qpt):
163            """
164            Build input file to compute GKQ.nc file from GS SCF input.
165            The calculation requires GS wavefunctions WFK, WFQ, a DDB file and a DVDB file
166            """
167            return scf_inp.new_with_vars(
168                optdriver=7,
169                eph_task=-2,
170                nqpt=1,
171                qpt=qpt,
172                ddb_ngqpt=ngqpt,  # q-mesh associated to the DDB file.
173                prtphdos=0,
174            )
175
176        # Now we compute matrix elements fully ab-initio for each q-point.
177        eph_work = Work()
178
179        qseen = set()
180        for task in work_qpath.phonon_tasks:
181            qpt = tuple(task.input["qpt"])
182            if qpt in qseen: continue
183            qseen.add(qpt)
184            t = eph_work.register_eph_task(make_eph_input(scf_input, ngqpt, qpt), deps=task.deps)
185            t.add_deps({work_qmesh: "DDB", work_qpath: "DVDB"})
186
187        flow.register_work(eph_work)
188
189        # Here we build another work to compute the gkq matrix elements
190        # with interpolated potentials along the q-path.
191        # The potentials are interpolated using the input ngqpt q-mesh.
192        if test_ft_interpolation:
193            inteph_work = Work()
194            qseen = set()
195            for task in work_qpath.phonon_tasks:
196                qpt = tuple(task.input["qpt"])
197                if qpt in qseen: continue
198                qseen.add(qpt)
199                eph_inp = make_eph_input(scf_input, ngqpt, qpt)
200                # Note eph_use_ftinterp 1 to force the interpolation of the DFPT potentials with eph_task -2.
201                eph_inp["eph_use_ftinterp"] = 1
202                t = inteph_work.register_eph_task(eph_inp, deps=task.deps)
203                t.add_deps({work_qmesh: ["DDB", "DVDB"]})
204            flow.register_work(inteph_work)
205
206        return flow
207