Merge pull request #31 from SMTG-UCL/mp-dos

Bugfix: passing a Pymatgen CompleteDOS to galore.process_pdos()

CHANGELOG.rst

@@ -12,6 +12,7 @@ Versioning <http://semver.org/>`__. The changelog format is inspired by
 - Update the [vasp] extra to handle some compatibility breaks between
   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()

galore/__init__.py

@@ -22,16 +22,16 @@
 from __future__ import print_function
 
 import os.path
-from itertools import repeat
 from collections import OrderedDict
+from collections.abc import Sequence
 import logging
 
 from math import sqrt, log
 import numpy as np
 from scipy.interpolate import interp1d
 
 import galore.formats
-from galore.cross_sections import get_cross_sections, cross_sections_info
+from galore.cross_sections import cross_sections_info
 
 
 def auto_limits(data_1d, padding=0.05):
@@ -134,9 +134,9 @@ def process_pdos(input=['vasprun.xml'],
             Files for processing. Vasp output or space-separated files with
                 XXX_EL_YYY.EXT filename pattern where EL is the element label.
                 We recommend SYSTEM_EL_dos.dat. Alternatively, a
-                `pymatgen.electronic_structure.dos.CompleteDos` can be provided.
-                Spin channels indicated by an (up) or (down) suffix in file
-                header will be combined for each orbital type.
+                `pymatgen.electronic_structure.dos.CompleteDos` can be
+                provided. Spin channels indicated by an (up) or (down) suffix
+                in file header will be combined for each orbital type.
         **kwargs:
             See main command reference
 
@@ -148,16 +148,15 @@ def process_pdos(input=['vasprun.xml'],
                  'el2': {'energy': values, 's': values, ...}, ...}
 
     """
-
-    if type(input) is str:
+    if isinstance(input, str) or not isinstance(input, Sequence):
         input = [input]
 
     # Read files into dict, check for consistency
     energy_label = None
     pdos_data = OrderedDict()
     for pdos_file in input:
-        if (galore.formats.is_xml(pdos_file) or
-                galore.formats.is_complete_dos(pdos_file)):
+        if (galore.formats.is_complete_dos(pdos_file)
+                or galore.formats.is_xml(pdos_file)):
             pdos_data = galore.formats.read_vasprun_pdos(pdos_file)
             kwargs['units'] = 'eV'
             break
@@ -427,7 +426,7 @@ def apply_orbital_weights(pdos_data, cross_sections):
             else:
                 try:
                     cs = cross_sections[el][orbital]
-                except KeyError as error:
+                except KeyError:
                     logging.warning("Could not find cross-section data for " +
                                     "element {0}, ".format(el) +
                                     "orbital {0}. ".format(orbital) +

galore/formats.py

@@ -301,6 +301,7 @@ def read_pdos_txt(filename, abs_values=True):
         return data[[col for col in data.dtype.names
                      if col not in spin_down_orbs]]
 
+
 def read_doscar(filename="DOSCAR"):
     """Read an x, y series of frequencies and DOS from a VASP DOSCAR file
 
@@ -351,7 +352,6 @@ def read_vasprun(filename='vasprun.xml'):
     """
     try:
         from pymatgen.io.vasp.outputs import Vasprun
-        from pymatgen.electronic_structure.core import Spin
     except ImportError as e:
         e.msg = "pymatgen package neccessary to load vasprun files"
         raise
@@ -471,6 +471,7 @@ def read_gpaw_pdos(filename, npts=50001, width=1e-3, ref='vbm'):
 
     return pdos_data
 
+
 def read_vasprun_totaldos(filename='vasprun.xml'):
     """Read an x, y series of energies and DOS from a VASP vasprun.xml file
 
@@ -500,7 +501,8 @@ def read_vasprun_pdos(filename='vasprun.xml'):
     Pymatgen must be present on the system to use this method
 
     Args:
-        filename (str): Path to vasprun.xml file or pymatgen CompleteDos object.
+        filename (str or CompleteDos):
+            Path to vasprun.xml file or pymatgen CompleteDos object.
 
     Returns:
         pdos_data (np.ndarray): PDOS data formatted as nestled OrderedDict of:

test/test.py

@@ -248,7 +248,6 @@ def test_read_complete_dos(self):
         self.assertEqual(pdos['Mg']['s'][150], 0.053)
         self.assertEqual(pdos['O']['p'][189], 0.004)
 
-
 txt_test_string = """# Frequency  Value
 0.000000e+00 0.000000e+00
 1.000000e+00 5.000000e-03