galore 0.10.0__py3-none-any.whl

galore/__init__.py

@@ -0,0 +1,441 @@
+###############################################################################
+#                                                                             #
+# GALORE: Gaussian and Lorentzian broadening for simulated spectra            #
+#                                                                             #
+# Developed by Scanlon Materials Theory Group at University College London    #
+# (c) 2017                                                                    #
+#                                                                             #
+###############################################################################
+#                                                                             #
+# This file is part of Galore. Galore is free software: you can redistribute  #
+# it and/or modify it under the terms of the GNU General Public License as    #
+# published by the Free Software Foundation, either version 3 of the License, #
+# or (at your option) any later version.  This program is distributed in the  #
+# hope that it will be useful, but WITHOUT ANY WARRANTY; without even the     #
+# implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.    #
+# See the GNU General Public License for more details.  You should have       #
+# received a copy of the GNU General Public License along with this program.  #
+# If not, see <http://www.gnu.org/licenses/>.                                 #
+#                                                                             #
+###############################################################################
+
+from __future__ import print_function
+
+from collections import OrderedDict
+from collections.abc import Sequence
+import logging
+from math import sqrt, log
+import os.path
+
+import numpy as np
+from scipy.interpolate import interp1d
+
+import galore.formats
+from galore.cross_sections import get_cross_sections, cross_sections_info
+
+
+def auto_limits(data_1d, padding=0.05):
+    """Return limiting values outside data range
+
+    Args:
+        data_1d (iterable): Data to obtain range from
+        padding (float): Scale factor for padding relative to data range
+
+    Returns:
+        (2-tuple) (xmin, xmax)
+
+    """
+    xmin, xmax = min(data_1d), max(data_1d)
+    auto_xmax = xmax + padding * (xmax - xmin)
+    auto_xmin = xmin - padding * (xmax - xmin)
+    return auto_xmin, auto_xmax
+
+
+def process_1d_data(input=['vasprun.xml'],
+                    gaussian=None, lorentzian=None,
+                    sampling=1e-2,
+                    xmin=None, xmax=None,
+                    spikes=False,
+                    **kwargs):
+    """Read 1D data series from files, process for output
+
+    Args:
+        input (str or 1-list):
+            Input data file. Pass as either a string or a list containing one
+            string
+        **kwargs:
+            See main command reference
+
+    Returns:
+        2-tuple (np.ndarray, np.ndarray):
+            Resampled x-values and corresponding broadened data as 1D numpy
+            arrays
+
+    """
+
+    if type(input) == str:
+        pass
+    elif len(input) > 1:
+        raise ValueError("Simple DOS only uses one input file, "
+                         "not list: {0}".format(input))
+    else:
+        input = input[0]
+
+    if not os.path.exists(input):
+        raise Exception(
+            "Input file {0} does not exist!".format(input))
+    if galore.formats.is_xml(input):
+        xy_data = galore.formats.read_vasprun_totaldos(input)
+    elif galore.formats.is_gpw(input):
+        xy_data = galore.formats.read_gpaw_totaldos(input)
+    elif galore.formats.is_doscar(input):
+        xy_data = galore.formats.read_doscar(input)
+    elif galore.formats.is_vasp_raman(input):
+        xy_data = galore.formats.read_vasp_raman(input)
+    elif galore.formats.is_csv(input):
+        xy_data = galore.formats.read_csv(input)
+    else:
+        xy_data = galore.formats.read_txt(input)
+
+    # Add 5% to data range if not specified
+    auto_xmin, auto_xmax = auto_limits(xy_data[:, 0], padding=0.05)
+    if xmax is None:
+        xmax = auto_xmax
+    if xmin is None:
+        xmin = auto_xmin
+
+    d = sampling
+    x_values = np.arange(xmin, xmax, d)
+    data_1d = galore.xy_to_1d(xy_data, x_values, spikes=spikes)
+
+    broadened_data = data_1d.copy()
+    if lorentzian:
+        broadened_data = galore.broaden(
+            broadened_data, d=d, dist='lorentzian', width=lorentzian)
+
+    if gaussian:
+        broadened_data = galore.broaden(
+            broadened_data, d=d, dist='gaussian', width=gaussian)
+
+    return (x_values, broadened_data)
+
+
+def process_pdos(input=['vasprun.xml'],
+                 gaussian=None, lorentzian=None,
+                 weighting=None, sampling=1e-2,
+                 xmin=None, xmax=None, flipx=False, **kwargs):
+    """Read PDOS from files, process for output
+
+    Args:
+        input (list or str):
+            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.
+        **kwargs:
+            See main command reference
+
+    Returns:
+        dict:
+            Weighted and resampled orbital data in format::
+
+                {'el1': {'energy': values, 's': values, 'p': values ...},
+                 'el2': {'energy': values, 's': values, ...}, ...}
+
+    """
+    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_complete_dos(pdos_file)
+                or galore.formats.is_xml(pdos_file)):
+            pdos_data = galore.formats.read_vasprun_pdos(pdos_file)
+            kwargs['units'] = 'eV'
+            break
+
+        elif galore.formats.is_gpw(pdos_file):
+            pdos_data = galore.formats.read_gpaw_pdos(pdos_file)
+            kwargs['units'] = 'eV'
+            break
+
+        if not os.path.exists(pdos_file):
+            raise Exception("Input file {0} does not "
+                            "exist!".format(input))
+
+        basename = os.path.basename(pdos_file)
+        try:
+            element = basename.split("_")[-2]
+        except IndexError:
+            raise Exception("Couldn't guess element name from filename. "
+                            "Please format filename as XXX_EL_YYY.EXT"
+                            "Where EL is the element label, and XXX, YYY "
+                            "and EXT are labels of your choice. We recommend"
+                            "SYSTEM_EL_dos.dat")
+
+        data = galore.formats.read_pdos_txt(pdos_file)
+
+        if energy_label is None:
+            energy_label = data.dtype.names[0]
+        else:
+            try:
+                assert data.dtype.names[0] == energy_label
+            except AssertionError as error:
+                error.args += ("Energy labels are not consistent "
+                               "between input files",)
+                raise
+
+        orbital_labels = data.dtype.names[1:]
+        pdos_data[element] = OrderedDict([('energy', data[energy_label])])
+        pdos_data[element].update(OrderedDict((orbital, data[orbital])
+                                  for orbital in orbital_labels))
+
+    # Work out sampling details; 5% pad added to data if no limits specified
+    # In x-flip mode, the user specifies these as binding energies so values
+    # are reversed while treating DOS data.
+    d = sampling
+    limits = (auto_limits(data['energy'], padding=0.05)
+              for (element, data) in pdos_data.items())
+    xmins, xmaxes = zip(*limits)
+
+    if xmax is None:
+        xmax = max(xmaxes)
+
+    if xmin is None:
+        xmin = min(xmins)
+
+    if flipx:
+        xmin, xmax = -xmax, -xmin
+
+    x_values = np.arange(xmin, xmax, d)
+
+    # Resample data into new dictionary
+    pdos_plotting_data = OrderedDict()
+    for element, el_data in pdos_data.items():
+        pdos_plotting_data[element] = OrderedDict([('energy', x_values)])
+        for orbital, orb_data in el_data.items():
+            if orbital == 'energy':
+                continue
+            elif orb_data is None:
+                continue
+
+            xy_data = np.column_stack([el_data['energy'], orb_data])
+
+            pdos_resampled = galore.xy_to_1d(xy_data, x_values)
+            broadened_data = pdos_resampled.copy()
+
+            if lorentzian:
+                broadened_data = galore.broaden(broadened_data, d=d,
+                                                dist='lorentzian',
+                                                width=lorentzian)
+
+            if gaussian:
+                broadened_data = galore.broaden(broadened_data, d=d,
+                                                dist='gaussian',
+                                                width=gaussian)
+
+            pdos_plotting_data[element][orbital] = broadened_data
+
+    if weighting:
+        cross_sections = galore.get_cross_sections(weighting,
+                                                   elements=pdos_data.keys())
+        pdos_plotting_data = galore.apply_orbital_weights(
+            pdos_plotting_data, cross_sections)
+
+    return pdos_plotting_data
+
+
+def xy_to_1d(xy, x_values, spikes=False):
+    """Convert a set of x,y coordinates to 1D array
+
+    Data is resampled to a given sequence of regularly-spaced x-values. By
+    default linear interpolation is used between neighbouring points, which is
+    suitable for resampling distribution data.
+
+    Where the data consists of a set of discrete energy levels, it should be
+    resampled to a series of "spikes". y-values are placed on the nearest
+    x-value by subtracting d/2 and rounding up. d is determined by examining
+    the first two elements of x_values.
+
+    Args:
+        xy: (ndarray) 2D numpy array of x, y values
+        x_values: (iterable) An evenly-spaced x-value mesh
+    Returns:
+        (np.array): re-sampled y values corresponding to x_values
+    """
+
+    x_values = np.array(x_values)
+    n_x_values = x_values.size
+    d = x_values[1] - x_values[0]
+
+    # Structured arrays are allowed, in which case first field is x,
+    # second is y. A bit of hackery is needed to slice these interchangeably.
+    if xy.dtype.names is None:
+        x_field, y_field = (Ellipsis, 0), (Ellipsis, 1)
+    else:
+        x_field, y_field = xy.dtype.names
+
+    if spikes:
+        spikes = np.zeros(n_x_values)
+        spike_locations = x_values.searchsorted(xy[x_field] - (0.5 * d))
+
+        for location, value in zip(spike_locations, xy[y_field]):
+            if location == 0 or location == n_x_values:
+                pass
+            else:
+                spikes[location] += value
+
+        return spikes
+
+    else:
+        y_func = interp1d(xy[x_field], xy[y_field],
+                          assume_sorted=False,
+                          bounds_error=False, fill_value=0)
+        return y_func(x_values)
+
+
+def delta(f1, f2, w=1):
+    """Compare two frequencies, return 1 if close"""
+    if abs(f1 - f2) <= 0.5 * w:
+        return 1
+    else:
+        return 0
+
+
+def lorentzian(f, f0=0, fwhm=1):
+    """Lorentzian function with height 1 centered on f0.
+
+    Args:
+        f (np.array): 1D array of x-values (e.g. frequencies)
+        f0 (float): Origin of function
+        fwhm (float): full-width half-maximum (FWHM);
+            i.e. the width of the function at half its maximum value.
+
+    """
+    return 0.5 * fwhm / (np.pi * ((f - f0)**2 + (0.5 * fwhm)**2))
+
+
+def gaussian(f, f0=0, fwhm=1):
+    """Gaussian function with height 1 centered on f0
+
+       f (np.array): 1D array of x-values (e.g. frequencies)
+       f0 (float): Origin of function
+       fwhm (float): full-width half-maximum (FWHM); i.e. the width of the
+       function at half its maximum value.
+
+    """
+    c = fwhm / (2 * sqrt(2 * log(2)))
+    return np.exp(-np.power(f - f0, 2) / (2 * c**2))
+
+
+def broaden(data, dist='lorentz', width=2, pad=False, d=1):
+    """Given a 1d data set, use convolution to apply a broadening function
+
+    Args:
+        data (np.array): 1D array of data points to broaden
+        dist (str): Type of distribution used for broadening. Currently only
+            "Lorentz" is supported.
+        width (float): Width parameter for broadening function. Determines the
+            full-width at half-maximum (FWHM) of the broadening function.
+        pad (float): Distance sampled on each side of broadening function.
+        d (float): x-axis distance associated with each sample in 1D data
+
+    """
+
+    if not pad:
+        pad = width * 20
+
+    if dist.lower() in ('lorentz', 'lorentzian'):
+        fwhm = width
+        broadening = lorentzian(np.arange(-pad, pad, d), f0=0, fwhm=fwhm)
+    elif dist.lower() in ('gauss', 'gaussian'):
+        fwhm = width
+        broadening = gaussian(np.arange(-pad, pad, d), f0=0, fwhm=fwhm)
+    else:
+        raise Exception('Broadening distribution '
+                        ' "{0}" not known.'.format(dist))
+
+    pad_points = int(pad / d)
+    broadened_data = np.convolve(broadening, data)
+    broadened_data = broadened_data[pad_points:len(data) + pad_points]
+
+    return broadened_data
+
+
+def apply_orbital_weights(pdos_data, cross_sections):
+    """Weight orbital intensities by cross-section for photoemission simulation
+
+    Args:
+        pdos_data (dict): DOS data in format::
+
+                {'el1': {'energy': values, 's': values, 'p': values ...},
+                 'el2': {'energy': values, 's': values, ...}, ...}
+
+             where DOS values are 1D numpy arrays. Orbital labels must match
+             cross_sections data. It is recommended to use
+             collections.OrderedDict instead of regular dictionaries, to ensure
+             consistent output.
+
+             In addition, the fields "citation", "link" and "energy" are
+             recognised and logged as "INFO".
+
+        cross_sections (dict): Weightings in format::
+
+                {'el1': {'1s': x1, '2s': x2, '2p': x3 ...},
+                 'el2': {'3s': y1, '3p': y2 ...}, ...}
+
+             The labels should correspond to the headers in the input data. It
+             is fine not so specify the level (e.g. use 's', 'p', etc.) as is
+             done in the sample data; however, this means that all levels are
+             being treated equally and hence probably the core levels will be
+             weighted incorrectly. It is possible to set the cross-section of
+             undesired orbitals (e.g. projection onto d-orbital for early
+             elements) to None; in this case the orbital will be dropped from
+             the returned data set.
+
+    Returns:
+        weighted_pdos_data (dict): Weighted data in same format as input
+    """
+
+    if type(cross_sections) != dict:
+        raise TypeError('Cross-section data should be a dictionary. Try using '
+                        'galore.get_cross_sections for a suitable data set.')
+
+    logging.info("Applying cross-section weighting values to PDOS:")
+
+    cross_sections_info(cross_sections, logging=logging)
+
+    logging.info("  Orbital cross-section weights per electron:")
+
+    weighted_pdos_data = OrderedDict()
+    for el, orbitals in pdos_data.items():
+        if 'warning' in cross_sections[el]:
+            logging.warning("  {0}: {1}".format(
+                el, cross_sections[el]['warning']))
+        weighted_orbitals = OrderedDict()
+        for orbital, data in orbitals.items():
+            if orbital == 'energy':
+                weighted_orbitals.update({'energy': data})
+            else:
+                try:
+                    cs = cross_sections[el][orbital]
+                except KeyError:
+                    logging.warning("Could not find cross-section data for " +
+                                    "element {0}, ".format(el) +
+                                    "orbital {0}. ".format(orbital) +
+                                    "Skipping this orbital.")
+                    cs = None
+
+                if cs is None:
+                    pass
+                else:
+                    logging.info("   {0} {1}: {2:.3e}".format(el, orbital, cs))
+                    weighted_orbitals.update({orbital: data * cs})
+
+        weighted_pdos_data.update({el: weighted_orbitals})
+
+    return weighted_pdos_data