Skip to content

Commit

Permalink
Added tests for summed_spin_channels options
Browse files Browse the repository at this point in the history
  • Loading branch information
JaGeo authored and JaGeo committed Feb 24, 2021
1 parent 5b21fa2 commit 54173a7
Show file tree
Hide file tree
Showing 4 changed files with 245 additions and 79 deletions.
101 changes: 38 additions & 63 deletions pymatgen/electronic_structure/cohp.py
Expand Up @@ -339,45 +339,32 @@ def get_cohp_by_label(self, label, summed_spin_channels=False):
Returns:
Returns the COHP object to simplify plotting
"""
# TODO: possibly add option to only get one spin channel as Spin.down in the future!
if label.lower() == "average":
divided_cohp = self.cohp
divided_icohp = self.icohp

if not summed_spin_channels:
if label.lower() == "average":
return Cohp(
efermi=self.efermi,
energies=self.energies,
cohp=self.cohp,
are_coops=self.are_coops,
icohp=self.icohp,
)
return Cohp(
efermi=self.efermi,
energies=self.energies,
cohp=self.all_cohps[label].get_cohp(spin=None, integrated=False),
are_coops=self.are_coops,
icohp=self.all_cohps[label].get_icohp(spin=None),
)
else:
divided_cohp = self.all_cohps[label].get_cohp(spin=None, integrated=False)
divided_icohp = self.all_cohps[label].get_icohp(spin=None)

if summed_spin_channels and Spin.down in self.cohp:
final_cohp = {}
final_icohp = {}
if label.lower() == "average":

divided_cohp = self.cohp
divided_icohp = self.icohp

else:
divided_cohp = self.all_cohps[label].get_cohp(spin=None, integrated=False)
divided_icohp = self.all_cohps[label].get_icohp(spin=None, integrated=False)
final_cohp[Spin.up] = np.sum([divided_cohp[Spin.up], divided_cohp[Spin.down]], axis=0)
final_icohp[Spin.up] = np.sum([divided_icohp[Spin.up], divided_icohp[Spin.down]], axis=0)
return Cohp(
efermi=self.efermi,
energies=self.energies,
cohp=final_cohp,
are_coops=self.are_coops,
icohp=self.all_cohps[label].get_icohp(spin=None),
)
else:
final_cohp = divided_cohp
final_icohp = divided_icohp

return Cohp(
efermi=self.efermi,
energies=self.energies,
cohp=final_cohp,
are_coops=self.are_coops,
icohp=final_icohp,
)

# TODO: add tests for summed_spin_channels (What happens if data is not spinpolarized?
def get_summed_cohp_by_label_list(self, label_list, divisor=1, summed_spin_channels=False):
"""
Returns a COHP object that includes a summed COHP divided by divisor
Expand Down Expand Up @@ -416,7 +403,6 @@ def get_summed_cohp_by_label_list(self, label_list, divisor=1, summed_spin_chann
if summed_spin_channels and Spin.down in summed_cohp:
final_cohp = {}
final_icohp = {}

final_cohp[Spin.up] = np.sum([divided_cohp[Spin.up], divided_cohp[Spin.down]], axis=0)
final_icohp[Spin.up] = np.sum([divided_icohp[Spin.up], divided_icohp[Spin.down]], axis=0)
else:
Expand Down Expand Up @@ -446,7 +432,6 @@ def get_summed_cohp_by_label_and_orbital_list(self, label_list, orbital_list, di
Returns:
Returns a COHP object including a summed COHP
"""
# TODO: add tests for summed_spin_channels
# check length of label_list and orbital_list:
if not len(label_list) == len(orbital_list):
raise ValueError("label_list and orbital_list don't have the same length!")
Expand Down Expand Up @@ -490,7 +475,6 @@ def get_summed_cohp_by_label_and_orbital_list(self, label_list, orbital_list, di
)

def get_orbital_resolved_cohp(self, label, orbitals, summed_spin_channels=False):
# TODO: add test summed_spin_channels
"""
Get orbital-resolved COHP.
Expand All @@ -500,6 +484,7 @@ def get_orbital_resolved_cohp(self, label, orbitals, summed_spin_channels=False)
orbitals: The orbitals as a label, or list or tuple of the form
[(n1, orbital1), (n2, orbital2)]. Orbitals can either be str,
int, or Orbital.
summed_spin_channels: bool, will sum the spin channels and return the sum in Spin.up if true
Returns:
Expand Down Expand Up @@ -535,36 +520,26 @@ def get_orbital_resolved_cohp(self, label, orbitals, summed_spin_channels=False)
except KeyError:
icohp = None

if not summed_spin_channels:
return Cohp(
self.efermi,
self.energies,
self.orb_res_cohp[label][orb_label]["COHP"],
icohp=icohp,
are_coops=self.are_coops,
)
else:
# TODO simplify code!
divided_cohp = self.orb_res_cohp[label][orb_label]["COHP"]
divided_icohp = self.orb_res_cohp[label][orb_label]["ICOHP"]
start_cohp = self.orb_res_cohp[label][orb_label]["COHP"]
start_icohp = icohp

if summed_spin_channels and Spin.down in start_cohp:
final_cohp = {}
final_icohp = {}
if Spin.down in divided_cohp:
final_cohp[Spin.up] = np.sum([divided_cohp[Spin.up], divided_cohp[Spin.down]], axis=0)
else:
final_cohp = divided_cohp
if Spin.down in divided_icohp:
final_icohp[Spin.up] = np.sum([divided_icohp[Spin.up], divided_icohp[Spin.down]], axis=0)
else:
final_icohp = divided_icohp

return Cohp(
self.efermi,
self.energies,
final_cohp,
icohp=final_icohp,
are_coops=self.are_coops,
)
final_cohp[Spin.up] = np.sum([start_cohp[Spin.up], start_cohp[Spin.down]], axis=0)
if start_icohp is not None:
final_icohp[Spin.up] = np.sum([start_icohp[Spin.up], start_icohp[Spin.down]], axis=0)
else:
final_cohp = start_cohp
final_icohp = start_icohp

return Cohp(
self.efermi,
self.energies,
final_cohp,
icohp=final_icohp,
are_coops=self.are_coops,
)

@classmethod
def from_dict(cls, d):
Expand Down
181 changes: 181 additions & 0 deletions pymatgen/electronic_structure/tests/test_cohp.py
Expand Up @@ -765,6 +765,12 @@ def setUp(self):
structure = os.path.join(test_dir, "POSCAR.Na2UO4")
self.cohp_lobster_forb = CompleteCohp.from_file("lobster", filename=filepath, structure_file=structure)

# #TODO: spinpolarized case:
filepath = os.path.join(test_dir, "environments", "COHPCAR.lobster.mp-190.gz")
structure = os.path.join(test_dir, "environments", "POSCAR.mp_190")
self.cohp_lobster_spin_polarized = CompleteCohp.from_file("lobster", filename=filepath,
structure_file=structure)

def test_attiributes(self):
self.assertFalse(self.cohp_lobster.are_coops)
self.assertFalse(self.cohp_lobster_dict.are_coops)
Expand Down Expand Up @@ -860,6 +866,39 @@ def test_get_cohp_by_label(self):
self.assertEqual(self.cohp_orb.get_icohp()[Spin.up][3], 0.0)
self.assertEqual(self.cohp_orb.get_cohp()[Spin.up][3], 0.0)

def test_get_cohp_by_label_summed_spin(self):
# files without spin polarization
self.assertAlmostEqual(self.cohp_orb.get_cohp_by_label("1", summed_spin_channels=True).energies[0], -11.7225)
self.assertAlmostEqual(self.cohp_orb.get_cohp_by_label("1", summed_spin_channels=True).energies[5], -11.47187)
self.assertFalse(self.cohp_orb.get_cohp_by_label("1", summed_spin_channels=True).are_coops)
self.assertAlmostEqual(self.cohp_orb.get_cohp_by_label("1", summed_spin_channels=True).cohp[Spin.up][0], 0.0)
self.assertAlmostEqual(self.cohp_orb.get_cohp_by_label("1", summed_spin_channels=True).cohp[Spin.up][300],
0.03392)
self.assertAlmostEqual(self.cohp_orb.get_cohp_by_label("average", summed_spin_channels=True).cohp[Spin.up][230],
-0.08792)
self.assertAlmostEqual(
self.cohp_orb.get_cohp_by_label("average", summed_spin_channels=True).energies[230],
-0.19368000000000007,
)
self.assertFalse(self.cohp_orb.get_cohp_by_label("average", summed_spin_channels=True).are_coops)

# file with spin polarization
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_cohp_by_label("1", summed_spin_channels=False).cohp[Spin.up][
300] * 2,
self.cohp_lobster_spin_polarized.get_cohp_by_label("1", summed_spin_channels=True).cohp[Spin.up][300])
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_cohp_by_label("1", summed_spin_channels=False).cohp[Spin.down][
300] * 2,
self.cohp_lobster_spin_polarized.get_cohp_by_label("1", summed_spin_channels=True).cohp[Spin.up][300])
self.assertAlmostEqual(self.cohp_lobster_spin_polarized.get_cohp_by_label("1",
summed_spin_channels=True).energies[
0], -15.03759 + 1.96204)
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_cohp_by_label("1", summed_spin_channels=True).energies[
5], -14.78697 + 1.96204)
self.assertFalse(self.cohp_lobster_spin_polarized.get_cohp_by_label("1", summed_spin_channels=True).are_coops)

def test_get_summed_cohp_by_label_list(self):
self.assertEqual(self.cohp_orb.get_summed_cohp_by_label_list(["1"]).energies[0], -11.7225)
self.assertEqual(
Expand All @@ -882,6 +921,55 @@ def test_get_summed_cohp_by_label_list(self):
0.03392,
)

def test_get_summed_cohp_by_label_list_summed_spin(self):
# files without spin polarization
self.assertEqual(self.cohp_orb.get_summed_cohp_by_label_list(["1"], summed_spin_channels=True).energies[0],
-11.7225)
self.assertEqual(
self.cohp_orb.get_summed_cohp_by_label_list(["1", "1"], summed_spin_channels=True).energies[0],
-11.7225,
)
self.assertEqual(self.cohp_orb.get_summed_cohp_by_label_list(["1"], summed_spin_channels=True).energies[5],
-11.47187)
self.assertFalse(self.cohp_orb.get_summed_cohp_by_label_list(["1"], summed_spin_channels=True).are_coops)
self.assertEqual(self.cohp_orb.get_summed_cohp_by_label_list(["1"], summed_spin_channels=True).cohp[Spin.up][0],
0.0)
self.assertEqual(
self.cohp_orb.get_summed_cohp_by_label_list(["1", "1"], summed_spin_channels=True).cohp[Spin.up][0],
0.0,
)
self.assertEqual(
self.cohp_orb.get_summed_cohp_by_label_list(["1", "1"], summed_spin_channels=True).cohp[Spin.up][300],
0.03392 * 2.0,
)
self.assertEqual(
self.cohp_orb.get_summed_cohp_by_label_list(["1", "1"], summed_spin_channels=True, divisor=2).cohp[
Spin.up][300],
0.03392,
)

# file with spin polarization
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_list(["1"], summed_spin_channels=False).cohp[
Spin.up][
300] * 2,
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_list(["1"], summed_spin_channels=True).cohp[
Spin.up][300])
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_list(["1"], summed_spin_channels=False).cohp[
Spin.down][
300] * 2,
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_list(["1"], summed_spin_channels=True).cohp[
Spin.up][300])
self.assertAlmostEqual(self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_list(["1", "1"],
summed_spin_channels=True).energies[
0], -15.03759 + 1.96204)
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_list(["1"], summed_spin_channels=True).energies[
5], -14.78697 + 1.96204)
self.assertFalse(self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_list(["1"],
summed_spin_channels=True).are_coops)

def test_get_summed_cohp_by_label_and_orbital_list(self):
ref = self.cohp_orb.orb_res_cohp["1"]["4s-4px"]
ref2 = self.cohp_orb.orb_res_cohp["1"]["4px-4pz"]
Expand All @@ -904,6 +992,60 @@ def test_get_summed_cohp_by_label_and_orbital_list(self):
with self.assertRaises(ValueError):
self.cohp_orb.get_summed_cohp_by_label_and_orbital_list(["1", "2"], ["4s-4px"])

def test_get_summed_cohp_by_label_and_orbital_list_summed_spin_channels(self):
ref = self.cohp_orb.orb_res_cohp["1"]["4s-4px"]
ref2 = self.cohp_orb.orb_res_cohp["1"]["4px-4pz"]
cohp_label = self.cohp_orb.get_summed_cohp_by_label_and_orbital_list(["1"], ["4s-4px"],
summed_spin_channels=True)
cohp_label2 = self.cohp_orb.get_summed_cohp_by_label_and_orbital_list(["1", "1"], ["4s-4px", "4s-4px"],
summed_spin_channels=True)
cohp_label2x = self.cohp_orb.get_summed_cohp_by_label_and_orbital_list(
["1", "1"], ["4s-4px", "4s-4px"], divisor=2, summed_spin_channels=True
)
cohp_label3 = self.cohp_orb.get_summed_cohp_by_label_and_orbital_list(["1", "1"], ["4px-4pz", "4s-4px"],
summed_spin_channels=True)

self.assertArrayEqual(cohp_label.cohp[Spin.up], ref["COHP"][Spin.up])
self.assertArrayEqual(cohp_label2.cohp[Spin.up], ref["COHP"][Spin.up] * 2.0)
self.assertArrayEqual(cohp_label3.cohp[Spin.up], ref["COHP"][Spin.up] + ref2["COHP"][Spin.up])
self.assertArrayEqual(cohp_label.icohp[Spin.up], ref["ICOHP"][Spin.up])
self.assertArrayEqual(cohp_label2.icohp[Spin.up], ref["ICOHP"][Spin.up] * 2.0)
self.assertArrayEqual(cohp_label2x.icohp[Spin.up], ref["ICOHP"][Spin.up])
self.assertArrayEqual(cohp_label3.icohp[Spin.up], ref["ICOHP"][Spin.up] + ref2["ICOHP"][Spin.up])
with self.assertRaises(ValueError):
self.cohp_orb.get_summed_cohp_by_label_and_orbital_list(["1"], ["4px-4pz", "4s-4px"],
summed_spin_channels=True)
with self.assertRaises(ValueError):
self.cohp_orb.get_summed_cohp_by_label_and_orbital_list(["1", "2"], ["4s-4px"], summed_spin_channels=True)

# files with spin polarization
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_and_orbital_list(["1"], ["6s-6s"],
summed_spin_channels=False).cohp[
Spin.up][
300] * 2,
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_and_orbital_list(["1"], ["6s-6s"],
summed_spin_channels=True).cohp[
Spin.up][300])
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_and_orbital_list(["1"], ["6s-6s"],
summed_spin_channels=False).cohp[
Spin.down][
300] * 2,
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_and_orbital_list(["1"], ["6s-6s"],
summed_spin_channels=True).cohp[
Spin.up][300])
self.assertAlmostEqual(self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_and_orbital_list(["1"],
["6s-6s"],
summed_spin_channels=True).energies[
0], -15.03759 + 1.96204)
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_and_orbital_list(["1"], ["6s-6s"],
summed_spin_channels=True).energies[
5], -14.78697 + 1.96204)
self.assertFalse(self.cohp_lobster_spin_polarized.get_summed_cohp_by_label_and_orbital_list(["1"], ["6s-6s"],
summed_spin_channels=True).are_coops)

def test_orbital_resolved_cohp(self):
# When read from a COHPCAR file, total COHPs are calculated from
# the orbital-resolved COHPs if the total is missing. This may be
Expand Down Expand Up @@ -934,6 +1076,45 @@ def test_orbital_resolved_cohp(self):
for cohp in cohps:
self.assertEqual(cohp.as_dict(), cohp_label.as_dict())

def test_orbital_resolved_cohp_summed_spin_channels(self):
ref = self.cohp_orb.orb_res_cohp["1"]["4s-4px"]
cohp_label = self.cohp_orb.get_orbital_resolved_cohp("1", "4s-4px", summed_spin_channels=True)
self.assertEqual(cohp_label.cohp, ref["COHP"])
self.assertEqual(cohp_label.icohp, ref["ICOHP"])
orbitals = [[Orbital.s, Orbital.px], ["s", "px"], [0, 3]]
cohps = [self.cohp_orb.get_orbital_resolved_cohp("1", [[4, orb[0]], [4, orb[1]]], summed_spin_channels=True) for
orb in orbitals]

for cohp in cohps:
self.assertEqual(cohp.as_dict(), cohp_label.as_dict())

# spin polarization
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_orbital_resolved_cohp("1", "6s-6s",
summed_spin_channels=False).cohp[
Spin.up][
300] * 2,
self.cohp_lobster_spin_polarized.get_orbital_resolved_cohp("1", "6s-6s",
summed_spin_channels=True).cohp[
Spin.up][300])
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_orbital_resolved_cohp("1", "6s-6s",
summed_spin_channels=False).cohp[
Spin.down][
300] * 2,
self.cohp_lobster_spin_polarized.get_orbital_resolved_cohp("1", "6s-6s",
summed_spin_channels=True).cohp[
Spin.up][300])
self.assertAlmostEqual(self.cohp_lobster_spin_polarized.get_orbital_resolved_cohp("1", "6s-6s",
summed_spin_channels=True).energies[
0], -15.03759 + 1.96204)
self.assertAlmostEqual(
self.cohp_lobster_spin_polarized.get_orbital_resolved_cohp("1", "6s-6s",
summed_spin_channels=True).energies[
5], -14.78697 + 1.96204)
self.assertFalse(self.cohp_lobster_spin_polarized.get_orbital_resolved_cohp("1", "6s-6s",
summed_spin_channels=True).are_coops)


if __name__ == "__main__":
unittest.main()

0 comments on commit 54173a7

Please sign in to comment.