class CohpPlotterTest(PymatgenTest):
def setUp(self):
path = os.path.join(test_dir, "cohp", "complete_cohp_lobster.json")
with open(os.path.join(path), "r") as f:
self.cohp = CompleteCohp.from_dict(json.load(f))
path = os.path.join(test_dir, "cohp", "complete_coop_lobster.json")
with open(os.path.join(path), "r") as f:
self.coop = CompleteCohp.from_dict(json.load(f))
self.cohp_plot = CohpPlotter(zero_at_efermi=False)
self.coop_plot = CohpPlotter(are_coops=True)
warnings.simplefilter("ignore")
def tearDown(self):
warnings.resetwarnings()
def test_attributes(self):
self.assertFalse(self.cohp_plot.are_coops)
self.assertTrue(self.coop_plot.are_coops)
self.assertFalse(self.cohp_plot.zero_at_efermi)
self.assertTrue(self.coop_plot.zero_at_efermi)
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
cohp_energies = self.cohp_plot._cohps["Fe8-Fe7"]["energies"]
self.assertEqual(len(cohp_energies), 301)
self.assertAlmostEqual(cohp_energies[0], -0.27768)
self.assertAlmostEqual(cohp_energies[-1], 14.77248)
self.coop_plot.add_cohp_dict(self.coop.all_cohps)
coop_energies = self.coop_plot._cohps["Bi5-Bi6"]["energies"]
self.assertEqual(len(coop_energies), 241)
self.assertAlmostEqual(coop_energies[0], -6.02510)
self.assertAlmostEqual(coop_energies[-1], 6.02510)
def test_add_cohp_dict(self):
# Sorts the populations by z-coordinates of the sites
def sortkeys(sites):
return sites[0].z, sites[1].z
sorted_keys = ["Bi2-Se8", "Bi2-Se9", "Bi4-Se9", "Bi4-Se12",
"Bi5-Se12", "Bi5-Bi6", "Bi6-Se11", "Bi3-Se11",
"Bi3-Se10", "Bi1-Se10", "Bi1-Se7"]
d_coop = self.coop_plot.get_cohp_dict()
self.assertEqual(len(d_coop), 0)
bonds = self.coop.bonds
self.coop_plot.add_cohp_dict(self.coop.all_cohps,
key_sort_func=lambda x:
sortkeys(bonds[x]["sites"]))
d_coop = self.coop_plot.get_cohp_dict()
self.assertEqual(len(d_coop), 11)
self.assertEqual(list(self.coop_plot._cohps.keys()), sorted_keys)
def test_get_cohp_dict(self):
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
d_cohp = self.cohp_plot.get_cohp_dict()
for bond in ["Fe8-Fe7", "Fe8-Fe9"]:
self.assertIn(bond, d_cohp)
def test_get_plot(self):
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
plt_cohp = self.cohp_plot.get_plot()
ax_cohp = plt_cohp.gca()
self.assertEqual(ax_cohp.get_xlabel(), "-COHP")
self.assertEqual(ax_cohp.get_ylabel(), "$E$ (eV)")
legend_labels = ax_cohp.get_legend_handles_labels()[1]
self.assertEqual(len(self.cohp_plot._cohps), len(legend_labels))
self.assertEqual(ax_cohp.lines[0].get_linestyle(), "-")
self.assertEqual(ax_cohp.lines[1].get_linestyle(), "--")
for label in legend_labels:
self.assertIn(label, self.cohp_plot._cohps)
linesindex = legend_labels.index("Fe8-Fe7")
linestyles = {Spin.up: '-', Spin.down: '--'}
cohp_fe_fe = self.cohp.all_cohps["Fe8-Fe7"]
for s, spin in enumerate([Spin.up, Spin.down]):
lines = ax_cohp.lines[2*linesindex+s]
self.assertArrayAlmostEqual(lines.get_xdata(),
-cohp_fe_fe.cohp[spin])
self.assertArrayAlmostEqual(lines.get_ydata(), self.cohp.energies)
self.assertEqual(lines.get_linestyle(), linestyles[spin])
plt_cohp.close()
plt_cohp = self.cohp_plot.get_plot(invert_axes=False,
plot_negative=False)
ax_cohp = plt_cohp.gca()
self.assertEqual(ax_cohp.get_xlabel(), "$E$ (eV)")
self.assertEqual(ax_cohp.get_ylabel(), "COHP")
for s, spin in enumerate([Spin.up, Spin.down]):
lines = ax_cohp.lines[2*linesindex+s]
self.assertArrayAlmostEqual(lines.get_xdata(), self.cohp.energies)
self.assertArrayAlmostEqual(lines.get_ydata(),
cohp_fe_fe.cohp[spin])
plt_cohp.close()
plt_cohp = self.cohp_plot.get_plot(integrated=True)
ax_cohp = plt_cohp.gca()
self.assertEqual(ax_cohp.get_xlabel(), "-ICOHP (eV)")
for s, spin in enumerate([Spin.up, Spin.down]):
lines = ax_cohp.lines[2*linesindex+s]
self.assertArrayAlmostEqual(lines.get_xdata(),
-cohp_fe_fe.icohp[spin])
coop_dict = {"Bi5-Bi6": self.coop.all_cohps["Bi5-Bi6"]}
self.coop_plot.add_cohp_dict(coop_dict)
plt_coop = self.coop_plot.get_plot()
ax_coop = plt_coop.gca()
self.assertEqual(ax_coop.get_xlabel(), "COOP")
self.assertEqual(ax_coop.get_ylabel(), "$E - E_f$ (eV)")
lines_coop = ax_coop.get_lines()[0]
self.assertArrayAlmostEqual(lines_coop.get_ydata(),
self.coop.energies - self.coop.efermi)
coop_bi_bi = self.coop.all_cohps["Bi5-Bi6"].cohp[Spin.up]
self.assertArrayAlmostEqual(lines_coop.get_xdata(), coop_bi_bi)
# Cleanup.
plt_cohp.close()
plt_coop.close()
def test_save_plot(self):
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
plt_cohp = self.cohp_plot.get_plot()
self.cohp_plot.save_plot("cohpplot.png")
self.assertTrue(os.path.isfile("cohpplot.png"))
os.remove("cohpplot.png")
plt_cohp.close()
class CohpPlotterTest(PymatgenTest):
def setUp(self):
path = os.path.join(test_dir, "cohp", "complete_cohp_lobster.json")
with open(os.path.join(path), "r") as f:
self.cohp = CompleteCohp.from_dict(json.load(f))
path = os.path.join(test_dir, "cohp", "complete_coop_lobster.json")
with open(os.path.join(path), "r") as f:
self.coop = CompleteCohp.from_dict(json.load(f))
self.cohp_plot = CohpPlotter(zero_at_efermi=False)
self.coop_plot = CohpPlotter(are_coops=True)
warnings.simplefilter("ignore")
def tearDown(self):
warnings.resetwarnings()
def test_attributes(self):
self.assertFalse(self.cohp_plot.are_coops)
self.assertTrue(self.coop_plot.are_coops)
self.assertFalse(self.cohp_plot.zero_at_efermi)
self.assertTrue(self.coop_plot.zero_at_efermi)
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
cohp_energies = self.cohp_plot._cohps["1"]["energies"]
self.assertEqual(len(cohp_energies), 301)
self.assertAlmostEqual(cohp_energies[0], -0.27768)
self.assertAlmostEqual(cohp_energies[-1], 14.77248)
self.coop_plot.add_cohp_dict(self.coop.all_cohps)
coop_energies = self.coop_plot._cohps["10"]["energies"]
self.assertEqual(len(coop_energies), 241)
self.assertAlmostEqual(coop_energies[0], -6.02510)
self.assertAlmostEqual(coop_energies[-1], 6.02510)
def test_add_cohp_dict(self):
# Sorts the populations by z-coordinates of the sites
def sortkeys(sites):
return sites[0].z, sites[1].z
sorted_keys = ["3", "4", "7", "8", "9", "10", "11", "6", "5", "2", "1"]
d_coop = self.coop_plot.get_cohp_dict()
self.assertEqual(len(d_coop), 0)
bonds = self.coop.bonds
self.coop_plot.add_cohp_dict(
self.coop.all_cohps,
key_sort_func=lambda x: sortkeys(bonds[x]["sites"]))
d_coop = self.coop_plot.get_cohp_dict()
self.assertEqual(len(d_coop), 11)
self.assertEqual(list(self.coop_plot._cohps.keys()), sorted_keys)
def test_get_cohp_dict(self):
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
d_cohp = self.cohp_plot.get_cohp_dict()
for bond in ["1", "2"]:
self.assertIn(bond, d_cohp)
def test_get_plot(self):
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
plt_cohp = self.cohp_plot.get_plot()
ax_cohp = plt_cohp.gca()
self.assertEqual(ax_cohp.get_xlabel(), "-COHP")
self.assertEqual(ax_cohp.get_ylabel(), "$E$ (eV)")
legend_labels = ax_cohp.get_legend_handles_labels()[1]
self.assertEqual(len(self.cohp_plot._cohps), len(legend_labels))
self.assertEqual(ax_cohp.lines[0].get_linestyle(), "-")
self.assertEqual(ax_cohp.lines[1].get_linestyle(), "--")
for label in legend_labels:
self.assertIn(label, self.cohp_plot._cohps)
linesindex = legend_labels.index("1")
linestyles = {Spin.up: '-', Spin.down: '--'}
cohp_fe_fe = self.cohp.all_cohps["1"]
for s, spin in enumerate([Spin.up, Spin.down]):
lines = ax_cohp.lines[2 * linesindex + s]
self.assertArrayAlmostEqual(lines.get_xdata(),
-cohp_fe_fe.cohp[spin])
self.assertArrayAlmostEqual(lines.get_ydata(), self.cohp.energies)
self.assertEqual(lines.get_linestyle(), linestyles[spin])
plt_cohp.close()
plt_cohp = self.cohp_plot.get_plot(invert_axes=False,
plot_negative=False)
ax_cohp = plt_cohp.gca()
self.assertEqual(ax_cohp.get_xlabel(), "$E$ (eV)")
self.assertEqual(ax_cohp.get_ylabel(), "COHP")
for s, spin in enumerate([Spin.up, Spin.down]):
lines = ax_cohp.lines[2 * linesindex + s]
self.assertArrayAlmostEqual(lines.get_xdata(), self.cohp.energies)
self.assertArrayAlmostEqual(lines.get_ydata(),
cohp_fe_fe.cohp[spin])
plt_cohp.close()
plt_cohp = self.cohp_plot.get_plot(integrated=True)
ax_cohp = plt_cohp.gca()
self.assertEqual(ax_cohp.get_xlabel(), "-ICOHP (eV)")
for s, spin in enumerate([Spin.up, Spin.down]):
lines = ax_cohp.lines[2 * linesindex + s]
self.assertArrayAlmostEqual(lines.get_xdata(),
-cohp_fe_fe.icohp[spin])
coop_dict = {"Bi5-Bi6": self.coop.all_cohps["10"]}
self.coop_plot.add_cohp_dict(coop_dict)
plt_coop = self.coop_plot.get_plot()
ax_coop = plt_coop.gca()
self.assertEqual(ax_coop.get_xlabel(), "COOP")
self.assertEqual(ax_coop.get_ylabel(), "$E - E_f$ (eV)")
lines_coop = ax_coop.get_lines()[0]
self.assertArrayAlmostEqual(lines_coop.get_ydata(),
self.coop.energies - self.coop.efermi)
coop_bi_bi = self.coop.all_cohps["10"].cohp[Spin.up]
self.assertArrayAlmostEqual(lines_coop.get_xdata(), coop_bi_bi)
# Cleanup.
plt_cohp.close()
plt_coop.close()
def test_save_plot(self):
self.cohp_plot.add_cohp_dict(self.cohp.all_cohps)
plt_cohp = self.cohp_plot.get_plot()
self.cohp_plot.save_plot("cohpplot.png")
self.assertTrue(os.path.isfile("cohpplot.png"))
os.remove("cohpplot.png")
plt_cohp.close()
def get_cohp(self, xlim=[-4, 4], ylim=[-10, 6]):
cohp = CompleteCohp.from_file('LOBSTER', filename=os.path.join(self.path, "COHPCAR.lobster"),
structure_file=os.path.join(self.path, "CONTCAR"))
plotter = CohpPlotter()
plotter.add_cohp(self.name, cohp)
plotter.save_plot(f'{self.path}.png', xlim=xlim, ylim=ylim, img_format='png')