1"""Test for GSR module"""
2import os
3import numpy as np
4import abipy.data as abidata
5import abipy.core
6import abipy.core.abinit_units as abu
7
8from pprint import pprint
9from abipy.core.testing import AbipyTest
10from abipy.electrons.gsr import GsrReader, GsrFile
11
12
13class GSRReaderTestCase(AbipyTest):
14
15    def test_read_Si2(self):
16        """Test the reading of GSR file."""
17        path = abidata.ref_file("si_scf_GSR.nc")
18
19        ref_dims = {
20            "number_of_spins": 1
21        }
22
23        ref_int_values = {
24            "space_group": 227,
25        }
26
27        ref_float_values = {
28            "etotal": -8.8652767680604807,
29        #    "primitive_vectors": np.reshape([0, 5.125, 5.125, 5.125, 0, 5.125,
30        #                                     5.125, 5.125, 0], (3,3)),
31        }
32
33        with GsrReader(path) as r:
34            assert r.ngroups == 1
35            varnames = r.read_varnames()
36            assert varnames and "ecut" in varnames
37
38            # Test dimensions.
39            for dimname, int_ref in ref_dims.items():
40                value = r.read_dimvalue(dimname)
41                self.assert_equal(value, int_ref)
42
43            # Test int variables
44            for varname, int_ref in ref_int_values.items():
45                value = r.read_value(varname)
46                self.assert_equal(value, int_ref)
47
48            # Test float variables
49            for varname, float_ref in ref_float_values.items():
50                value = r.read_value(varname)
51                self.assert_almost_equal(value, float_ref)
52
53            # Reading non-existent variables or dims should raise a subclass of NetcdReder.
54            with self.assertRaises(GsrReader.Error): r.read_value("foobar")
55            with self.assertRaises(GsrReader.Error): r.read_dimvalue("foobar")
56
57            r.print_tree()
58            for group in r.walk_tree():
59                print("group: " + str(group))
60
61            # Initialize pymatgen structure from GSR.
62            structure = r.read_structure()
63            assert isinstance(structure, abipy.core.Structure)
64
65
66class GSRFileTestCase(AbipyTest):
67
68    def test_gsr_silicon(self):
69        """spin unpolarized GSR file"""
70
71        with GsrFile(abidata.ref_file("si_scf_GSR.nc")) as gsr:
72            assert gsr.basename == "si_scf_GSR.nc"
73            assert gsr.relpath == os.path.relpath(abidata.ref_file("si_scf_GSR.nc"))
74            assert gsr.filetype
75            assert gsr.filestat()
76            assert len(gsr.ncdump())
77            repr(gsr); str(gsr)
78            assert gsr.to_string(verbose=2)
79            assert gsr.abinit_version == "8.0.6"
80            str(gsr.ebands)
81            assert gsr.filepath == abidata.ref_file("si_scf_GSR.nc")
82            assert gsr.nsppol == 1
83            assert gsr.mband == 8 and gsr.nband == 8 and gsr.nelect == 8 and len(gsr.kpoints) == 29
84            assert gsr.mband == gsr.hdr.mband
85            assert "nelect" in gsr.hdr and gsr.nelect == gsr.hdr.nelect
86            self.assert_almost_equal(gsr.energy.to("Ha"), -8.86527676798556)
87            self.assert_almost_equal(gsr.energy_per_atom * len(gsr.structure), gsr.energy)
88
89            assert gsr.params["nband"] == 8
90            assert gsr.params["nkpt"] == 29
91
92            # Test energy_terms
93            eterms = gsr.energy_terms
94            repr(eterms); str(eterms)
95            assert eterms.to_string(with_doc=True)
96            self.assert_almost_equal(eterms.e_xc.to("Ha"), -3.51815936301812)
97            self.assert_almost_equal(eterms.e_nonlocalpsp.to("Ha"), 1.91660690901782)
98            self.assert_almost_equal(eterms.e_kinetic.to("Ha"), 2.96421325671218)
99            self.assert_almost_equal(eterms.e_fermie.to("Ha"), 0.205739364929368)
100
101            # Forces and stress
102            self.assert_almost_equal(gsr.cart_forces.to("Ha bohr^-1").flat,
103               [-1.14726679671674e-28, -3.76037290483622e-29, 5.65937773808884e-29,
104                 1.14726679671674e-28, 3.76037290483622e-29, -5.65937773808884e-29])
105
106            self.assert_almost_equal(gsr.max_force, 0)
107            assert gsr.force_stats()
108            assert gsr.residm > 0
109            assert str(gsr.xc) == "LDA_XC_TETER93"
110
111            #self.assert_almost_equal(gsr.cart_stress_tensor.flat,
112            # Cartesian components of stress tensor (hartree/bohr^3)
113            #  sigma(1 1)=  1.77139311E-04  sigma(3 2)=  0.00000000E+00
114            #  sigma(2 2)=  1.77139311E-04  sigma(3 1)=  0.00000000E+00
115            #  sigma(3 3)=  1.77139311E-04  sigma(2 1)=  2.67294316E-15
116            for i in range(3):
117                self.assert_almost_equal(gsr.cart_stress_tensor[0, 0], 1.77139311E-04 * abu.HaBohr3_GPa)
118            self.assert_almost_equal(gsr.pressure, -5.211617575719521)
119
120            # Test pymatgen computed_entries
121            for inc_structure in (True, False):
122                e = gsr.get_computed_entry(inc_structure=inc_structure)
123                str(e)
124                d = e.as_dict()
125                if inc_structure: assert "structure" in d
126                assert d["energy"] == gsr.energy
127                assert gsr.energy == e.energy
128
129            if self.has_matplotlib():
130                assert gsr.plot_ebands(show=False)
131                assert gsr.plot_ebands_with_edos(edos=gsr.get_edos(), show=False)
132
133            if self.has_nbformat():
134                gsr.write_notebook(nbpath=self.get_tmpname(text=True))
135
136            if self.has_panel():
137                assert hasattr(gsr.get_panel(), "show")
138
139
140class GsrRobotTest(AbipyTest):
141
142    def test_gsr_robot(self):
143        """Testing GSR robot"""
144        from abipy import abilab
145        gsr_ibz = abidata.ref_file("si_scf_GSR.nc")
146        robot = abilab.GsrRobot()
147        robot.add_file("gsr0", gsr_ibz)
148        assert len(robot.abifiles) == 1
149        assert "gsr0" in robot.keys()
150        assert "gsr0" in robot.labels
151        assert robot.EXT == "GSR"
152        repr(robot); str(robot)
153        assert robot.to_string(verbose=2)
154
155	# Cannot have same label
156        with self.assertRaises(ValueError):
157            robot.add_file("gsr0", gsr_ibz)
158
159        assert len(robot) == 1 and not robot.exceptions
160        robot.add_file("gsr1", abilab.abiopen(gsr_ibz))
161        assert len(robot) == 2
162        robot.show_files()
163        assert not robot.has_different_structures()
164        with self.assertRaises(AttributeError):
165            robot.is_sortable("foobar", raise_exc=True)
166        assert not robot.is_sortable("foobar")
167        # Test different syntax.
168        assert robot.is_sortable("nkpt")         # gsr.nkpt
169        assert robot.is_sortable("ebands.nkpt")  # gsr.ebands.nkpt
170        assert robot.is_sortable("ecut")         # in gsr.params
171
172        dfs = robot.get_structure_dataframes()
173        assert dfs.lattice is not None
174        assert dfs.coords is not None
175        assert len(dfs.structures) == len(robot)
176
177        assert len(robot.get_ebands_plotter(kselect="path")) == 0
178        filter_abifile = lambda gsr: gsr.ebands.kpoints.is_ibz
179        assert len(robot.get_ebands_plotter(filter_abifile=filter_abifile)) == 2
180
181        ebands_plotter = robot.get_ebands_plotter()
182        edos_plotter = robot.get_edos_plotter()
183
184        if self.has_matplotlib():
185            assert ebands_plotter.gridplot(show=False)
186            assert robot.combiplot_ebands(show=False)
187            assert robot.gridplot_ebands(show=False)
188            assert robot.boxplot_ebands(show=False)
189            assert robot.combiboxplot_ebands(show=False)
190
191            assert edos_plotter.gridplot(show=False)
192            assert robot.combiplot_edos(show=False)
193            assert robot.gridplot_edos(show=False)
194
195            assert robot.plot_gsr_convergence(show=False)
196            assert robot.plot_gsr_convergence(sortby="nkpt", hue="tsmear", show=False)
197            y_vars = ["energy", "structure.lattice.a", "structure.volume"]
198            assert robot.plot_convergence_items(y_vars, sortby="nkpt", hue="tsmear", show=False)
199
200            assert robot.plot_egaps(show=False)
201            assert robot.plot_egaps(sortby="nkpt", hue="tsmear")
202            assert robot.gridplot_with_hue("structure.formula", show=False)
203
204        if self.has_panel():
205            assert hasattr(robot.get_panel(), "show")
206
207	# Get pandas dataframe.
208        df = robot.get_dataframe()
209        assert "energy" in df
210        self.assert_equal(df["ecut"].values, 6.0)
211        self.assert_almost_equal(df["energy"].values, -241.2364683, decimal=4)
212
213        df_params = robot.get_params_dataframe()
214        assert "nband" in df_params
215
216        assert "alpha" in robot.get_lattice_dataframe()
217        assert hasattr(robot.get_coords_dataframe(), "keys")
218
219        eterms_df = robot.get_energyterms_dataframe(iref=0)
220        assert "energy" in eterms_df
221        assert "ecut" in eterms_df
222        assert "nkpt" in eterms_df
223
224        if self.has_matplotlib():
225            assert robot.plot_xy_with_hue(df, x="nkpt", y="pressure", hue="a", show=False)
226
227        # Note: This is not a real EOS since we have a single volume.
228        # But testing is better than not testing.
229        r = robot.get_eos_fits_dataframe()
230        assert hasattr(r, "fits") and hasattr(r, "dataframe")
231
232        if self.has_matplotlib():
233            assert robot.gridplot_eos(show=False)
234
235        if self.has_nbformat():
236            robot.write_notebook(nbpath=self.get_tmpname(text=True))
237
238        robot.close()
239
240        # Test other class methods
241        filepath = abidata.ref_file("si_scf_GSR.nc")
242        robot = abilab.GsrRobot.from_dirs(os.path.dirname(filepath), abspath=True)
243        assert len(robot) == 2
244        assert robot[filepath].filepath == filepath
245
246        # Test from_dir_glob
247        pattern = "%s/*/si_ebands/" % abidata.dirpath
248        same_robot = abilab.GsrRobot.from_dir_glob(pattern, abspath=True)
249        assert len(same_robot) == 2
250        assert set(robot.labels) == set(same_robot.labels)
251
252        robot.close()
253        same_robot.close()
254