1"""Integration tests for GW flows."""
2
3import abipy.data as abidata
4import abipy.abilab as abilab
5import abipy.flowtk as flowtk
6
7#from abipy.core.testing import has_abinit, has_matplotlib
8
9
10def make_g0w0_inputs(ngkpt, tvars):
11    """
12    Input files for the calculation of the GW corrections.
13
14    Returns: gs_input, nscf_input, scr_input, sigma_input
15    """
16    multi = abilab.MultiDataset(structure=abidata.cif_file("si.cif"),
17                                pseudos=abidata.pseudos("14si.pspnc"), ndtset=4)
18
19    # This grid is the most economical, but does not contain the Gamma point.
20    scf_kmesh = dict(
21        ngkpt=ngkpt,
22        shiftk=[0.5, 0.5, 0.5,
23                0.5, 0.0, 0.0,
24                0.0, 0.5, 0.0,
25                0.0, 0.0, 0.5]
26    )
27
28    # This grid contains the Gamma point, which is the point at which
29    # we will compute the (direct) band gap.
30    gw_kmesh = dict(
31        ngkpt=ngkpt,
32        shiftk=[0.0, 0.0, 0.0,
33                0.0, 0.5, 0.5,
34                0.5, 0.0, 0.5,
35                0.5, 0.5, 0.0]
36    )
37
38    # Global variables. gw_kmesh is used in all datasets except DATASET 1.
39    ecut = 4
40
41    multi.set_vars(
42        ecut=ecut,
43        pawecutdg=ecut*2 if multi.ispaw else None,
44        istwfk="*1",
45        paral_kgb=tvars.paral_kgb,
46        gwpara=2,
47    )
48    multi.set_kmesh(**gw_kmesh)
49
50    # Dataset 1 (GS run)
51    multi[0].set_kmesh(**scf_kmesh)
52    multi[0].set_vars(tolvrs=1e-6, nband=4)
53
54    # Dataset 2 (NSCF run)
55    multi[1].set_vars(iscf=-2,
56                      tolwfr=1e-10,
57                      nband=10,
58                      nbdbuf=2)
59
60    # Dataset3: Calculation of the screening.
61    multi[2].set_vars(
62        optdriver=3,
63        nband=8,
64        ecutwfn=ecut,
65        symchi=1,
66        inclvkb=0,
67        ecuteps=2.0,
68    )
69
70    # Dataset4: Calculation of the Self-Energy matrix elements (GW corrections)
71    kptgw = [
72         -2.50000000E-01, -2.50000000E-01,  0.00000000E+00,
73         -2.50000000E-01,  2.50000000E-01,  0.00000000E+00,
74          5.00000000E-01,  5.00000000E-01,  0.00000000E+00,
75         -2.50000000E-01,  5.00000000E-01,  2.50000000E-01,
76          5.00000000E-01,  0.00000000E+00,  0.00000000E+00,
77          0.00000000E+00,  0.00000000E+00,  0.00000000E+00,
78      ]
79
80    multi[3].set_vars(
81            optdriver=4,
82            nband=10,
83            ecutwfn=ecut,
84            ecuteps=2.0,
85            ecutsigx=2.0,
86            symsigma=1,
87            #gw_qprange=0,
88    )
89
90    bdgw = [4, 5]
91    multi[3].set_kptgw(kptgw, bdgw)
92
93    return multi.split_datasets()
94
95
96def itest_g0w0_flow(fwp, tvars):
97    """Test flow for G0W0 calculations."""
98    scf, nscf, scr, sig = make_g0w0_inputs(ngkpt=[2, 2, 2], tvars=tvars)
99
100    flow = flowtk.g0w0_flow(fwp.workdir, scf, nscf, scr, sig, manager=fwp.manager)
101    # Will remove output files at run-time.
102    flow.set_garbage_collector()
103    flow.build_and_pickle_dump(abivalidate=True)
104
105    for task in flow[0]:
106        task.start_and_wait()
107
108    flow.check_status(show=True)
109    assert all(work.finalized for work in flow)
110    if not flow.all_ok:
111        flow.debug()
112        raise RuntimeError()
113
114    scf_task = flow[0][0]
115    nscf_task = flow[0][1]
116    scr_task = flow[0][2]
117    sig_task = flow[0][3]
118
119    # Test garbage_collector
120    # The WFK|SCR file should have been removed because we call set_garbage_collector
121    assert not scf_task.outdir.has_abiext("WFK")
122    assert not nscf_task.outdir.has_abiext("WFK")
123    assert not scr_task.outdir.has_abiext("SCR")
124    assert not scr_task.outdir.has_abiext("SUS")
125
126    # The sigma task should produce a SIGRES file.
127    sigfile = sig_task.outdir.list_filepaths(wildcard="*SIGRES.nc")[0]
128    assert sigfile
129    with abilab.abiopen(sigfile) as sigres:
130        sigres.to_string(verbose=2)
131        assert sigres.nsppol == 1
132
133    # Test SigmaTask inspect method
134    #if has_matplotlib():
135        #sig_task.inspect(show=False)
136
137    # Test get_results for Sigma and Scr
138    scr_task.get_results()
139    sig_task.get_results()
140
141    # Test SCR.nc file (this is optional)
142    if scr_task.scr_path:
143        with scr_task.open_scr() as scr:
144            scr.to_string(verbose=2)
145            assert len(scr.wpts) == 2
146            assert scr.nwre == 1 and scr.nwim == 1
147            for iq, qpoint in enumerate(scr.qpoints[:2]):
148                #print(qpoint)
149                qpt, iqcheck = scr.reader.find_qpoint_fileindex(qpoint)
150                assert iqcheck == iq
151                em1 = scr.get_em1(qpoint)
152                #print(em1)
153
154    # TODO Add more tests
155    #assert flow.validate_json_schema()
156
157
158def itest_g0w0qptdm_flow(fwp, tvars):
159    """Integration test for G0W0WithQptdmFlow."""
160    scf, nscf, scr, sig = make_g0w0_inputs(ngkpt=[2, 2, 2], tvars=tvars)
161
162    flow = flowtk.G0W0WithQptdmFlow(fwp.workdir, scf, nscf, scr, sig, manager=fwp.manager)
163
164    # Enable garbage collector at the flow level.
165    # Note that here we have tp use this policy because tasks are created dynamically
166    #flow.set_garbage_collector(policy="task")
167    flow.set_garbage_collector(policy="flow")
168
169    assert len(flow) == 3
170    bands_work = flow[0]
171    scr_work = flow[1]
172    sigma_work = flow[2]
173
174    assert scr_work.depends_on(bands_work.nscf_task)
175    assert not scr_work.depends_on(bands_work.scf_task)
176
177    for sigma_task in sigma_work:
178        #print("sigma_task.deps", sigma_task.deps)
179        assert sigma_task.depends_on(bands_work.nscf_task)
180        assert not sigma_task.depends_on(bands_work.scf_task)
181        assert sigma_task.depends_on(scr_work)
182
183    flow.build_and_pickle_dump(abivalidate=True)
184    flow.show_dependencies()
185    # This call is needed to connect the node and enable
186    # the callbacks, otherwise the scheduler enters a deadlock.
187    flow.connect_signals()
188
189    # Run the flow.
190    fwp.scheduler.add_flow(flow)
191    assert fwp.scheduler.start() == 0
192    assert not fwp.scheduler.exceptions
193
194    flow.show_status()
195    assert all(work.finalized for work in flow)
196    if not flow.all_ok:
197        flow.debug()
198        raise RuntimeError()
199
200    # Test set_garbage_collector
201    # The WFK|SCR file should have been removed because we call set_garbage_collector
202    #assert not scf_task.outdir.has_abiext("WFK")
203    #assert not nscf_task.outdir.has_abiext("WFK")
204    #assert not scr_task.outdir.has_abiext("SCR")
205    #assert not scr_task.outdir.has_abiext("SUS")
206
207    # The SCR file produced by scr_work should have been removed
208    assert not scr_work.outdir.has_abiext("SCR")
209
210    #assert flow.validate_json_schema()
211
212    flow.finalize()
213
214
215def itest_htc_g0w0(fwp, tvars):
216    """Testing G0W0Work."""
217    structure = abilab.Structure.from_file(abidata.cif_file("si.cif"))
218    pseudos = abidata.pseudos("14si.pspnc")
219
220    flow = flowtk.Flow(fwp.workdir, manager=fwp.manager)
221
222    scf_kppa = 10
223    nscf_nband = 10
224    #nscf_ngkpt = [4,4,4]
225    #nscf_shiftk = [0.0, 0.0, 0.0]
226    ecut, ecuteps, ecutsigx = 4, 2, 3
227    #scr_nband = 50
228    #sigma_nband = 50
229
230    extra_abivars = dict(
231        ecut=ecut,
232        istwfk="*1",
233        paral_kgb=tvars.paral_kgb,
234    )
235
236    multi = abilab.g0w0_with_ppmodel_inputs(
237        structure, pseudos,
238        scf_kppa, nscf_nband, ecuteps, ecutsigx,
239        ecut=ecut, pawecutdg=None,
240        accuracy="normal", spin_mode="unpolarized", smearing=None,
241        #ppmodel="godby", charge=0.0, scf_algorithm=None, inclvkb=2, scr_nband=None,
242        #sigma_nband=None, gw_qprange=1):
243
244    )
245    multi.set_vars(paral_kgb=tvars.paral_kgb)
246
247    scf_input, nscf_input, scr_input, sigma_input = multi.split_datasets()
248    work = flowtk.G0W0Work(scf_input, nscf_input, scr_input, sigma_input)
249
250    flow.register_work(work)
251    flow.allocate()
252    flow.connect_signals()
253
254    fwp.scheduler.add_flow(flow)
255    assert fwp.scheduler.start() == 0
256    assert not fwp.scheduler.exceptions
257    assert fwp.scheduler.nlaunch == 4
258
259    # The sigma task should produce a SCR file.
260    assert len(work[2].outdir.list_filepaths(wildcard="*SCR")) == 1
261
262    flow.show_status()
263    if not flow.all_ok:
264        flow.debug()
265        raise RuntimeError()
266
267    assert all(work.finalized for work in flow)
268
269    #assert flow.validate_json_schema()
270