1"""Tests for electrons.bse module"""
2import abipy.data as abidata
3
4from abipy.core.testing import AbipyTest
5from abipy.electrons.scissors import *
6from abipy import abilab
7
8
9class TestScissors(AbipyTest):
10
11    def test_scissors_polyfit(self):
12        """Testing scissors from SIGRES file."""
13        # Get the quasiparticle results from the SIGRES.nc database.
14        with abilab.abiopen(abidata.ref_file("si_g0w0ppm_nband30_SIGRES.nc")) as sigma_file:
15            qplist_spin = sigma_file.qplist_spin
16
17        # Construct the scissors operator
18        domains = [[-10, 6.1], [6.1, 18]]
19        scissors = qplist_spin[0].build_scissors(domains, bounds=None)
20        #scissors = qplist_spin[0].build_scissors(domains, bounds=None, plot=True)
21
22        # Read the KS band energies computed on the k-path
23        with abilab.abiopen(abidata.ref_file("si_nscf_GSR.nc")) as nc:
24           ks_bands = nc.ebands
25
26        # Read the KS band energies computed on the Monkhorst-Pack (MP) mesh
27        # and compute the DOS with the Gaussian method
28        with abilab.abiopen(abidata.ref_file("si_scf_GSR.nc")) as nc:
29            ks_mpbands = nc.ebands
30        ks_dos = ks_mpbands.get_edos()
31
32        # Apply the scissors operator first on the KS band structure
33        # along the k-path then on the energies computed with the MP mesh.
34        qp_bands = ks_bands.apply_scissors(scissors)
35        qp_mpbands = ks_mpbands.apply_scissors(scissors)
36
37        # Compute the DOS with the modified QPState energies.
38        qp_dos = qp_mpbands.get_edos()
39
40        # Plot the LDA and the QPState band structure with matplotlib.
41        if self.has_matplotlib():
42            plotter = abilab.ElectronBandsPlotter()
43            plotter.add_ebands("LDA", ks_bands, edos=ks_dos)
44            plotter.add_ebands("LDA+scissors(e)", qp_bands, edos=qp_dos)
45
46            # By default, the two band energies are shifted wrt to *their* fermi level.
47            # Use e=0 if you don't want to shift the eigenvalus
48            # so that it's possible to visualize the QP corrections.
49            plotter.combiplot(title="Silicon band structure", show=False)
50            plotter.gridplot(title="Silicon band structure", show=False)
51
52    def test_builder(self):
53        """Testing ScissorsBuilder."""
54        builder = ScissorsBuilder.from_file(abidata.ref_file("si_g0w0ppm_nband30_SIGRES.nc"))
55        assert builder.nsppol == 1
56
57        # To select the domains esplicitly (optional)
58        builder.build(domains_spin=[[-10, 6.02], [6.1, 20]])
59
60        # Test pickle.
61        #tmp_path = self.get_tmpname(suffix=".pickle")
62        #builder.pickle_dump(tmp_path)
63        #new = ScissorsBuilder.pickle_load(tmp_path)
64
65        if self.has_matplotlib():
66            # To plot the QP results as function of the KS energy:
67            assert builder.plot_qpe_vs_e0(show=False)
68            # To compare the fitted results with the ab-initio data:
69            assert builder.plot_fit(show=False)
70
71            # To plot the corrected bands:
72            bands_filepath = abidata.ref_file("si_nscf_WFK.nc")
73            assert builder.plot_qpbands(bands_filepath, show=False)
74
75            dos_filepath = abidata.ref_file("si_scf_GSR.nc")
76            assert builder.plot_qpbands(bands_filepath, dos_filepath=dos_filepath,
77                    dos_args=None, show=False)
78