1# coding: utf-8
2# Copyright (c) Pymatgen Development Team.
3# Distributed under the terms of the MIT License.
4
5import os
6import unittest
7
8from pymatgen.io.pwscf import PWInput, PWInputError, PWOutput
9from pymatgen.util.testing import PymatgenTest
10
11
12class PWInputTest(PymatgenTest):
13    def test_init(self):
14        s = self.get_structure("Li2O")
15        self.assertRaises(
16            PWInputError,
17            PWInput,
18            s,
19            control={"calculation": "scf", "pseudo_dir": "./"},
20            pseudo={"Li": "Li.pbe-n-kjpaw_psl.0.1.UPF"},
21        )
22
23    def test_str_mixed_oxidation(self):
24        s = self.get_structure("Li2O")
25        s.remove_oxidation_states()
26        s[1] = "Li1"
27        pw = PWInput(
28            s,
29            control={"calculation": "scf", "pseudo_dir": "./"},
30            pseudo={
31                "Li": "Li.pbe-n-kjpaw_psl.0.1.UPF",
32                "Li+": "Li.pbe-n-kjpaw_psl.0.1.UPF",
33                "O": "O.pbe-n-kjpaw_psl.0.1.UPF",
34            },
35            system={"ecutwfc": 50},
36        )
37        ans = """&CONTROL
38  calculation = 'scf',
39  pseudo_dir = './',
40/
41&SYSTEM
42  ecutwfc = 50,
43  ibrav = 0,
44  nat = 3,
45  ntyp = 3,
46/
47&ELECTRONS
48/
49&IONS
50/
51&CELL
52/
53ATOMIC_SPECIES
54  Li  6.9410 Li.pbe-n-kjpaw_psl.0.1.UPF
55  Li+  6.9410 Li.pbe-n-kjpaw_psl.0.1.UPF
56  O  15.9994 O.pbe-n-kjpaw_psl.0.1.UPF
57ATOMIC_POSITIONS crystal
58  O 0.000000 0.000000 0.000000
59  Li+ 0.750178 0.750178 0.750178
60  Li 0.249822 0.249822 0.249822
61K_POINTS automatic
62  1 1 1 0 0 0
63CELL_PARAMETERS angstrom
64  2.917389 0.097894 1.520005
65  0.964634 2.755036 1.520005
66  0.133206 0.097894 3.286918
67"""
68        self.assertEqual(pw.__str__().strip(), ans.strip())
69
70    def test_str_without_oxidation(self):
71        s = self.get_structure("Li2O")
72        s.remove_oxidation_states()
73        pw = PWInput(
74            s,
75            control={"calculation": "scf", "pseudo_dir": "./"},
76            pseudo={
77                "Li": "Li.pbe-n-kjpaw_psl.0.1.UPF",
78                "O": "O.pbe-n-kjpaw_psl.0.1.UPF",
79            },
80            system={"ecutwfc": 50},
81        )
82        ans = """&CONTROL
83  calculation = 'scf',
84  pseudo_dir = './',
85/
86&SYSTEM
87  ecutwfc = 50,
88  ibrav = 0,
89  nat = 3,
90  ntyp = 2,
91/
92&ELECTRONS
93/
94&IONS
95/
96&CELL
97/
98ATOMIC_SPECIES
99  Li  6.9410 Li.pbe-n-kjpaw_psl.0.1.UPF
100  O  15.9994 O.pbe-n-kjpaw_psl.0.1.UPF
101ATOMIC_POSITIONS crystal
102  O 0.000000 0.000000 0.000000
103  Li 0.750178 0.750178 0.750178
104  Li 0.249822 0.249822 0.249822
105K_POINTS automatic
106  1 1 1 0 0 0
107CELL_PARAMETERS angstrom
108  2.917389 0.097894 1.520005
109  0.964634 2.755036 1.520005
110  0.133206 0.097894 3.286918
111"""
112        self.assertEqual(pw.__str__().strip(), ans.strip())
113
114    def test_str_with_oxidation(self):
115        s = self.get_structure("Li2O")
116
117        pw = PWInput(
118            s,
119            control={"calculation": "scf", "pseudo_dir": "./"},
120            pseudo={
121                "Li+": "Li.pbe-n-kjpaw_psl.0.1.UPF",
122                "O2-": "O.pbe-n-kjpaw_psl.0.1.UPF",
123            },
124            system={"ecutwfc": 50},
125        )
126        ans = """&CONTROL
127  calculation = 'scf',
128  pseudo_dir = './',
129/
130&SYSTEM
131  ecutwfc = 50,
132  ibrav = 0,
133  nat = 3,
134  ntyp = 2,
135/
136&ELECTRONS
137/
138&IONS
139/
140&CELL
141/
142ATOMIC_SPECIES
143  Li+  6.9410 Li.pbe-n-kjpaw_psl.0.1.UPF
144  O2-  15.9994 O.pbe-n-kjpaw_psl.0.1.UPF
145ATOMIC_POSITIONS crystal
146  O2- 0.000000 0.000000 0.000000
147  Li+ 0.750178 0.750178 0.750178
148  Li+ 0.249822 0.249822 0.249822
149K_POINTS automatic
150  1 1 1 0 0 0
151CELL_PARAMETERS angstrom
152  2.917389 0.097894 1.520005
153  0.964634 2.755036 1.520005
154  0.133206 0.097894 3.286918
155"""
156        self.assertEqual(pw.__str__().strip(), ans.strip())
157
158
159class PWOuputTest(PymatgenTest):
160    def setUp(self):
161        self.pwout = PWOutput(os.path.join(PymatgenTest.TEST_FILES_DIR, "Si.pwscf.out"))
162
163    def test_properties(self):
164        self.assertAlmostEqual(self.pwout.final_energy, -93.45259708)
165
166    def test_get_celldm(self):
167        self.assertAlmostEqual(self.pwout.get_celldm(1), 10.323)
168        for i in range(2, 7):
169            self.assertAlmostEqual(self.pwout.get_celldm(i), 0)
170
171
172if __name__ == "__main__":
173    unittest.main()
174