plotSpectra.py

#!/usr/bin/env python3

"""

plotSpectra
===========

This module contains functions which are useful to post-process spectrum data.
Also functionality for plotting spectra generated by the Quanty software.

"""

import numpy as np
import sys
from math import pi

from .constants import k_B


def fermi_smearing(x, y, fwhm):
    r"""
    Return Fermi smeared values.

    Smearing is done by convolution of the input data
    with the function:

    .. math:: f(x) = -\frac{\partial n(x)}{\partial x},

    where

    .. math:: n(x) = \frac{1}{\exp{(\beta x)}+1}

    is the Fermi-dirac distribution.

    Assumes uniform grid.

    Parameters
    ----------
    x : array
    y : array
    fwhm : float
        Full Width Half Max.

    """
    # Inverse temperature
    beta = 4.*np.log(1+np.sqrt(2))/fwhm
    # Grid spacing (assume uniform grid)
    dx = x[1] - x[0]
    # Grid spacing should be smaller than the
    # standard deviation
    assert dx < fwhm
    # Create a mesh for the smearing function
    xg = np.arange(-3*fwhm, 3*fwhm, dx)
    # Create smear function
    g = beta*np.exp(beta*xg)/(np.exp(beta*xg)+1)**2
    # This check ensures the return array has the
    # same shape as x
    assert len(x) > len(xg)
    # Convolute y with smearing function
    smeared = dx*np.convolve(y, g, mode='same')
    return smeared


def gaussian_smearing(x, y, fwhm):
    """
    Return Gaussian smeared values of variable y.

    Assumes uniform grid.

    Parameters
    ----------
    x - array
    y - array
    fwhm - float
        Full Width Half Max.

    """
    # the standard deviation of the Gaussian
    std = 1./(2*np.sqrt(2*np.log(2)))*fwhm
    # grid spacing (assume uniform grid)
    dx = x[1]-x[0]
    # Grid spacing should be smaller than the
    # standard deviation
    assert dx < std
    # Create a mesh for the Gaussian function
    xg = np.arange(-5*std, 5*std, dx)
    # Create a gaussian centered around zero
    g = 1./np.sqrt(2*pi*std**2)*np.exp(-xg**2/(2*std**2))
    # This check ensures the return array has the
    # same shape as x
    assert len(x) > len(xg)
    # Convolute y with Gaussian
    smeared = dx*np.convolve(y, g, mode='same')
    return smeared


def lorentzian_smearing(x, y, fwhm):
    """
    Return Lorenzian smeared values of variable y.

    Assumes uniform grid.

    Parameters
    ----------
    x - array
    y - array
    fwhm - float
        Full Width Half Max.

    """
    # delta variable in the Lorentzian
    d = fwhm/2.
    # Grid spacing (assume uniform grid)
    dx = x[1]-x[0]
    # Grid spacing should be smaller than the
    # standard deviation
    assert dx < d
    # Create a mesh for the Lorentzian function
    xg = np.arange(-10*d, 10*d, dx)
    # Create a Lorentzian centered around zero
    g = 1/pi*d/(xg**2+d**2)
    # This check ensures the return array has the
    # same shape as x
    assert len(x) > len(xg)
    # Convolute y with Lorentzian
    smeared = dx*np.convolve(y, g, mode='same')
    return smeared


# Reading and processing XA spectra, generated by Quanty

def read_Quanty_data_file(filename):
    """
    Return the spectra read from file.

    """
    with open(filename) as f:
        content = f.readlines()
    save_line = False
    x = []
    for line in content:
        columns = line.split()
        if save_line:
            x.append([float(item) for item in columns])
        if columns[0] == 'Energy':
            save_line = True
    return np.array(x)


def read_Quanty_output_file(filename):
    """
    Read the eigenstate information read from file.

    """
    with open(filename) as f:
        content = f.readlines()
    save_line = False
    x = []
    for line in content:
        columns = line.split()
        if save_line and columns:
            try:
                x.append([float(e) for e in columns])
            except ValueError:
                break
        if len(columns) >= 2 and columns[1] == 'states':
            save_line = True
    return np.array(x)


def get_normalization(x, minus_intensity, loc, h, peak=0,
                      peakorder='topdown'):
    y = -minus_intensity
    i = find_peak_index(x, y, peak, peakorder)
    shift = loc - x[i]
    scale = h/minus_intensity[i]
    return shift, scale


def find_peak_index(x, y, peak_nbr=0, peakorder='topdown'):
    r"""
    Return position index of the `peak_nbr` highest or
    leftest peak.

    Parameters
    ----------
    x : (N) array
        Energy mesh.
    y : (N) array
    peak_nbr : int
        Peak index.
    peakorder : str
        'topdown' or 'leftright'
        How to sort peaks in variable `y`.

    """
    if peakorder == 'topdown' and peak_nbr == 0:
        return np.argmax(y)  # this is a special case
    else:
        mask1 = np.logical_and(0 <= np.diff(y[:-1]),
                               np.diff(y[1:]) < 0)
        mask2 = np.logical_and(0 < np.diff(y[:-1]),
                               np.diff(y[1:]) <= 0)
        mask = np.logical_or(mask1, mask2)
        peak_x = x[1:-1][mask]
        peak_y = y[1:-1][mask]
        # print peak_x
        if peakorder == 'topdown':
            indices = np.argsort(peak_y)
            index = indices[-1-peak_nbr]
        elif peakorder == 'leftright':
            indices = np.argsort(peak_x)
            index = indices[peak_nbr]
        else:
            sys.exit('Warning: Value of peakorder variable'
                     'is incorrect.')
        return np.argmin(np.abs(x-peak_x[index]))


def thermal_average(energies, observable, T=300):
    """
    Returns thermally averaged observables.

    Assumes all relevant states are included.
    Thus, no not check e.g. if the Boltzmann weight
    of the last state is small.

    Parameters
    ----------
    energies : list(N)
        energies[i] is the energy of state i.
    observable : list(N,...)
        observable[i,...] are observables for state i.
    T : float
        Temperature

    """
    if len(energies) != np.shape(observable)[0]:
        raise ValueError("Passed array is not of the right shape")
    z = 0
    e_average = 0
    o_average = 0
    weights = np.zeros(np.shape(energies), dtype=np.float)
    shift = np.min(energies)
    for j, (e, o) in enumerate(zip(energies, observable)):
        weight = np.exp(-(e-shift)/(k_B*T))
        z += weight
        weights[j] = weight
        e_average += weight*e
        o_average += weight*o
    e_average /= z
    o_average /= z
    weights /= z
    return o_average


def get_index_unique(x, xtol=0.001):
    """
    Return (first) indices of non degenerate values and corresponding degeneracy.

    """
    ind = []
    degen = []
    for i, e in enumerate(x):
        if i == 0:
            ind.append(i)
            degen.append(1)
        elif e-x[i-1] > xtol:
            ind.append(i)
            degen.append(1)
        else:
            degen[-1] += 1
    return ind, degen


def get_Quanty_spectrum(folder='.', xas_file='XASSpec.dat',
                        output_file='output.txt', T=300,
                        peak=None, loc=None, h=None,
                        peakorder='topdown', tol=4e-5):
    """
    Return the thermally averaged spectrum.

    It can also be shifted in energy and scaled
    in intensity by the parameters peak, loc and h.

    Checks if enough energies have been computed.
    The Boltzmann weight of the state with the highest
    energy should be smaller than variable tol, otherwise
    a warning is printed.
    """
    xas = read_Quanty_data_file(folder + '/' + xas_file)
    w = xas[:, 0]
    energies = read_Quanty_output_file(folder + '/' + output_file)[:, 1]

    if tol < np.exp(-(energies[-1]-energies[0])/(k_B*T)):
        print('Warning: Perhaps too few eigenenergies considered.')
        print('E-E0 =')
        print(energies-energies[0])

    xasa = thermal_average(energies, np.transpose(xas[:, 2::2]), T)
    if peak is None and loc is None and h is None:
        return w, xasa
    elif peak is None or loc is None or h is None:
        print('Either none or all of the parameters:'
              'peak, loc and h should be specified.')
        return w, xasa
    else:
        shift, scale = get_normalization(w, xasa, loc, h,
                                         peak=peak,
                                         peakorder=peakorder)
        return w + shift, scale*xasa