diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 614c3d6..b064cb0 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -13,6 +13,9 @@ Versioning `__. The changelog format is inspired by
dependencies and different Python versions.
- Fix some incorrect values in Al k-alpha XPS cross-sections
- BUGFIX: Pymatgen CompleteDOS was not correctly accepted by galore.process_pdos()
+- Implement previously ineffective "offset" option in
+ galore.plot.plot_pdos(), add a matching option to
+ galore.plot.plot_tdos()
`[0.6.1] `__ - 2018-11-19
-----------------------------------------------------------------------------------
diff --git a/galore/plot.py b/galore/plot.py
index a98fec9..36f58b2 100644
--- a/galore/plot.py
+++ b/galore/plot.py
@@ -97,7 +97,7 @@ def add_overlay(plt, overlay, overlay_scale=None, overlay_offset=0.,
def plot_pdos(pdos_data, ax=None, total=True, show_orbitals=True,
- offset=0, flipx=False, **kwargs):
+ offset=0., flipx=False, **kwargs):
"""Plot a projected density of states (PDOS)
Args:
@@ -141,9 +141,9 @@ def plot_pdos(pdos_data, ax=None, total=True, show_orbitals=True,
# Field 'energy' must be present, other fields are orbitals
assert 'energy' in el_data.keys()
if flipx:
- x_data = -el_data['energy']
+ x_data = -el_data['energy'] + offset
else:
- x_data = el_data['energy']
+ x_data = el_data['energy'] + offset
orbitals = list(el_data.keys())
orbitals.remove('energy')
@@ -183,7 +183,7 @@ def plot_pdos(pdos_data, ax=None, total=True, show_orbitals=True,
return plt
-def plot_tdos(xdata, ydata, ax=None, **kwargs):
+def plot_tdos(xdata, ydata, ax=None, offset=0., **kwargs):
"""Plot a total DOS (i.e. 1D dataset)
@@ -191,6 +191,7 @@ def plot_tdos(xdata, ydata, ax=None, **kwargs):
xdata (iterable): x-values (energy, frequency etc.)
ydata (iterable): Corresponding y-values (DOS or measurement intensity)
show (bool): Display plot
+ offset (float): Energy shift to x-axis
ax (matplotlib.Axes): If provided, plot onto existing Axes object. If
None, a new Figure will be created and the pyplot instance will be
returned.
@@ -212,7 +213,9 @@ def plot_tdos(xdata, ydata, ax=None, **kwargs):
ax = fig.add_subplot(1, 1, 1)
if kwargs['flipx']:
- xdata = -xdata
+ xdata = -xdata + offset
+ else:
+ xdata = xdata + offset
ax.plot(xdata, ydata, 'C0-')
ax.set_xlim([min(xdata), max(xdata)])
diff --git a/test/test_plot.py b/test/test_plot.py
index 6c645c8..cbba1ad 100644
--- a/test/test_plot.py
+++ b/test/test_plot.py
@@ -70,26 +70,29 @@ def test_plot_pdos(self):
('energy', np.array([1, 2, 3, 4, 5])),
('s', np.array([1, 1, 0, 0, 0])),
('p', np.array([0, 1, 2, 2, 1]))]))])
+ offset = 0.5
- galore.plot.plot_pdos(pdos_data, ax=ax)
+
+ galore.plot.plot_pdos(pdos_data, ax=ax, offset=offset)
line1 = ax.lines[0]
xy1 = line1.get_xydata()
- self.assertEqual(xy1[1, 0], 2)
+ self.assertEqual(xy1[1, 0], 2 + offset)
self.assertEqual(xy1[1, 1], 0)
line2 = ax.lines[1]
xy2 = line2.get_xydata()
- self.assertEqual(xy2[2, 0], 3)
+ self.assertEqual(xy2[2, 0], 3 + offset)
self.assertEqual(xy2[2, 1], 4)
line3 = ax.lines[2]
xy3 = line3.get_xydata()
- self.assertEqual(xy3[0, 0], 1)
+ self.assertEqual(xy3[0, 0], 1 + offset)
self.assertEqual(xy2[0, 1], 1)
tdos = ax.lines[4]
xyt = tdos.get_xydata()
+ self.assertEqual(xyt[2, 0], 3 + offset)
self.assertEqual(xyt[3, 1], 1 + 1 + 0 + 2)
@@ -99,11 +102,13 @@ def test_plot_tdos(self):
ax = fig.add_subplot(1, 1, 1)
xvals = np.linspace(-5, 5, 21)
- galore.plot.plot_tdos(xvals, xvals**2, ax=ax)
+ offset = 0.8
+
+ galore.plot.plot_tdos(xvals, xvals**2, ax=ax, offset=offset)
self.assertEqual(len(ax.lines), 1)
- self.assertEqual(ax.lines[0].get_xydata()[11, 0], 0.5)
- self.assertAlmostEqual(ax.lines[0].get_xydata()[18, 0], 4.0)
+ self.assertEqual(ax.lines[0].get_xydata()[11, 0], 0.5 + offset)
+ self.assertAlmostEqual(ax.lines[0].get_xydata()[18, 1], 4.0**2)
self.assertAlmostEqual(ax.get_ylim()[0], 0)
self.assertAlmostEqual(ax.get_ylim()[1], 1.1 * 5**2)