1"""Tests for wannier90 module"""
2import os
3import numpy as np
4import abipy.data as abidata
5
6from abipy import abilab
7from abipy.core.testing import AbipyTest
8
9
10class TestWoutFile(AbipyTest):
11
12    def test_example01_gaas(self):
13        """Parsing example01_gaas.wout"""
14        filepath = os.path.join(abidata.dirpath, "refs", "wannier90", "example01_gaas.wout")
15        with abilab.abiopen(filepath) as wout:
16            repr(wout); str(wout)
17            assert wout.to_string(verbose=2)
18            assert wout.version == "2.1.0+git"
19            assert wout.structure.formula == "Ga1 As1"
20            self.assert_almost_equal(wout.structure.frac_coords.ravel(),
21                    [0.00000, 0.00000, 0.00000, 0.25000, 0.25000, 0.25000])
22            assert not wout.warnings
23            for k in ("MAIN", "WANNIERISE"):
24                assert k in wout.params_section
25            assert not wout.use_disentangle
26            assert wout.nwan == 4
27            assert np.all(wout.grid_size == 2)
28            assert len(wout.conv_df) == 20 + 1
29            assert wout.conv_df.O_D[0] == 0.0083198 and wout.conv_df.O_OD[0] == 0.5036294
30            assert wout.conv_df.O_D[20] == 0.0080300 and wout.conv_df.O_OD[20] == 0.5019880
31
32            # numpy array (nwan, nstep, ...)
33            #natom = len(wout.structure)
34            nstep = 21
35            assert wout.wf_centers.shape == (wout.nwan, nstep, 3)
36            self.assert_equal(wout.wf_centers[1, 20], [-0.866253,  0.866253,  0.866253])
37            assert wout.wf_spreads.shape == (wout.nwan, nstep)
38            self.assert_equal(wout.wf_spreads[1, 20], 1.11672024)
39
40            if self.has_matplotlib():
41                assert wout.plot(show=False)
42                assert wout.plot_centers_spread(show=False)
43
44            if self.has_nbformat():
45                assert wout.write_notebook(nbpath=self.get_tmpname(text=True))
46
47    def test_example03_silicon(self):
48        """Parsing example02_silicon.wout with DISENTANGLE"""
49        filepath = os.path.join(abidata.dirpath, "refs", "wannier90", "example03_silicon.wout")
50        with abilab.abiopen(filepath) as wout:
51            repr(wout); str(wout)
52            assert wout.to_string(verbose=2)
53            assert wout.version == "2.1.0+git"
54            assert wout.structure.formula == "Si2"
55            self.assert_almost_equal(wout.structure.frac_coords.ravel(),
56                    [-0.25000, 0.75000, -0.25000, 0.00000, 0.00000, 0.00000])
57            assert not wout.warnings
58            for k in ("MAIN", "WANNIERISE", "DISENTANGLE"):
59                assert k in wout.params_section
60            assert wout.use_disentangle
61            assert wout.nwan == 8
62            assert np.all(wout.grid_size == 4)
63            assert len(wout.conv_df) == 6 + 1
64            assert wout.conv_df.O_D[1] == 0.1213986 and wout.conv_df.O_OD[1] == 2.7017701
65            assert wout.conv_df.O_D.values[-1] == 0.1054702 and wout.conv_df.O_OD.values[-1] == 2.5449106
66
67            # numpy array (nwan, nstep, ...)
68            nstep = len(wout.conv_df.O_D)
69            assert wout.wf_centers.shape == (wout.nwan, nstep, 3)
70            self.assert_equal(wout.wf_centers[7, -1], [0.888643,  0.888652,  1.810090 ])
71            assert wout.wf_spreads.shape == (wout.nwan, nstep)
72            self.assert_equal(wout.wf_spreads[7, -1], 1.81245236)
73
74            if self.has_matplotlib():
75                assert wout.plot(show=False)
76                assert wout.plot_centers_spread(show=False)
77
78            if self.has_nbformat():
79                assert wout.write_notebook(nbpath=self.get_tmpname(text=True))
80