Source code for sloth.math.convolution1D

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Generic discrete convolution

Description
-----------

 This is a manual (not optimized!) implementation of discrete 1D
 convolution intended for spectroscopy analysis. The difference with
 commonly used methods is the possibility to adapt the convolution
 kernel for each convolution point, e.g. change the FWHM of the
 Gaussian kernel as a function of the energy scale.

Resources
---------

.. [WPconv] <http://en.wikipedia.org/wiki/Convolution#Discrete_convolution>
.. [Fisher] <http://homepages.inf.ed.ac.uk/rbf/HIPR2/convolve.htm>
.. [GP1202] <http://glowingpython.blogspot.fr/2012/02/convolution-with-numpy.html>

TODO
----
- [] get_ene_index: substitute with more elegant
     'np.argmin(np.abs(ene-(x)))'
- [] atan_gamma_fdmnes: define atan_gamma as in FDMNES
"""
import os
import math
import subprocess
from optparse import OptionParser
from datetime import date
from string import Template
import numpy as np

MODNAME = "_math"
DEBUG = 0

# No deps on Larch: used only if you want to access this as Larch plugin
HAS_LARCH = False
try:
    from larch import use_plugin_path

    use_plugin_path("math")
    HAS_LARCH = True
except ImportError:
    pass

# <LINESHAPES> #


[docs] def gaussian(x, cen=0, sigma=1, fwhm=False, peak=None): """1 dimensional Gaussian function (https://en.wikipedia.org/wiki/Gaussian_function) Parameters ---------- x : array cen : [0] center, x0 sigma : [1] standard deviation, FWHM = 2*sqrt(2*ln(2)) * sigma =~ 2.35482 * sigma fwhm : [False] if True, the given sigma is assumed as fwhm and then converted accordingly peak : [None] if None, peak = 1 / math.sqrt(2*math.pi), the distribution integrate to 1 """ if fwhm is True: sigma = sigma / 2 * math.sqrt(2 * math.log(2)) if peak is None: peak = 1.0 / math.sqrt(2 * math.pi) return peak * np.exp(-((1.0 * x - cen) ** 2) / (2 * sigma ** 2))
[docs] def lorentzian(x, cen=0, gamma=1, peak=None): """1 dimensional Lorentzian Parameters ---------- x : array cen : [0] center, x0 gamma : [1] half width at half maximum peak : [None] if None, peak = 1 / (math.pi*sigma), the distribution integrate to 1 """ if peak is None: peak = 1.0 / (math.pi * gamma) return peak * (1.0 / (1.0 + ((1.0 * x - cen) / gamma) ** 2))
# </LINESHAPES> #
[docs] def get_ene_index(ene, cen, hwhm): """ returns the min/max indexes for array ene at (cen-hwhm) and (cen+hwhm) very similar to index_of in larch """ try: if (cen - hwhm) <= min(ene): ene_imin = 0 else: ene_imin = max(np.where(ene < (cen - hwhm))[0]) if (cen + hwhm) >= max(ene): ene_imax = len(ene) - 1 else: ene_imax = min(np.where(ene > (cen + hwhm))[0]) return ene_imin, ene_imax except Exception: print("index not found for {0} +/- {1}".format(cen, hwhm)) return None, None
[docs] def lin_gamma(ene, fwhm=1.0, linbroad=None): """returns constant or linear energy-dependent broadening Parameters ---------- ene : energy array in eV fwhm : first full width at half maximum in eV linbroad : list of 3-elements giving 'second full width at half maximum' 'energy starting point of the linear increase' 'energy ending point of the linear increase' """ w = np.ones_like(ene) if linbroad is None: return w * fwhm else: try: fwhm2 = linbroad[0] e1 = linbroad[1] e2 = linbroad[2] except Exception: raise ValueError("wrong format for linbroad") for en, ee in enumerate(ene): if ee < e1: w[en] *= fwhm elif ee <= e2: wlin = fwhm + (ee - e1) * (fwhm2 - fwhm) / (e2 - e1) w[en] *= wlin elif ee >= e2: w[en] *= fwhm2 return w
[docs] def atan_gamma(ene, gamma_hole, gamma_max=15.0, e0=0, eslope=1.0): """returns arctangent-like broadening, $\Gamma(E)$ ..math \Gamma(E) = \Gamma_{hole} + \Gamma_{max} * ( \arctan( \frac{E-E_{0}}{E_{slope}} ) / \pi + 1/2) ) """ if eslope == 0: print("Warning: eslope cannot be zero, using default value of 1") eslope = 1.0 return gamma_hole + gamma_max * ((np.arctan((ene - e0) / eslope) / np.pi) + 0.5)
[docs] def conv(e, mu, kernel="gaussian", fwhm_e=None, efermi=None): """ linear broadening Parameters ---------- e : x-axis (energy) mu : f(x) to convolve with g(x) kernel, mu(energy) kernel : convolution kernel, g(x) 'gaussian' 'lorentzian' fwhm_e: the full width half maximum in eV for the kernel broadening. It is an array of size 'e' with constants or an energy-dependent values determined by a function as 'lin_gamma()' or 'atan_gamma()' """ f = np.copy(mu) z = np.zeros_like(f) if efermi is not None: # ief = index_nearest(e, efermi) ief = np.argmin(np.abs(e - efermi)) f[0:ief] *= 0 if e.shape != fwhm_e.shape: print("Error: 'fwhm_e' does not have the same shape of 'e'") return 0 # linar fit upper part of the spectrum to avoid border effects # polyfit => pf lpf = int(len(e) / 2) cpf = np.polyfit(e[-lpf:], f[-lpf:], 1) fpf = np.poly1d(cpf) # extend upper energy border to 3*fhwm_e[-1] estep = e[-1] - e[-2] eup = np.append(e, np.arange(e[-1] + estep, e[-1] + 3 * fwhm_e[-1], estep)) for n in range(len(f)): # from now on I change e with eup eimin, eimax = get_ene_index(eup, eup[n], 1.5 * fwhm_e[n]) if (eimin is None) or (eimax is None): if DEBUG: raise IndexError("e[{0}]".format(n)) if len(range(eimin, eimax)) % 2 == 0: kx = eup[eimin:eimax + 1] # odd range centered at the convolution point else: kx = eup[eimin:eimax] # kernel ### hwhm = fwhm_e[n] / 2.0 if "gauss" in kernel.lower(): ky = gaussian(kx, cen=eup[n], sigma=hwhm) elif "lor" in kernel.lower(): ky = lorentzian(kx, cen=eup[n], gamma=hwhm) else: raise ValueError("convolution kernel '{0}' not implemented".format(kernel)) ky = ky / ky.sum() # normalize zn = 0 lk = int(len(kx)) for mf, mg in zip(range(-int(lk / 2), int(lk / 2) + 1), range(lk)): if ((n + mf) >= 0) and ((n + mf) < len(f)): zn += f[n + mf] * ky[mg] elif (n + mf) >= 0: zn += fpf(eup[n + mf]) * ky[mg] z[n] = zn return z
[docs] def glinbroad(e, mu, fwhm_e=None, efermi=None, _larch=None): """gaussian linear convolution in Larch """ if _larch is None: raise Warning("larch broken?") return conv(e, mu, kernel="gaussian", fwhm_e=fwhm_e, efermi=efermi)
glinbroad.__doc__ = conv.__doc__ # CONVOLUTION WITH FDMNES VIA SYSTEM CALL #
[docs] class FdmnesConv(object): """ Performs convolution with FDMNES within Python """ def __init__(self, opts=None, calcroot=None, fn_in=None, fn_out=None): if opts is None: self.opts = dict( creator="FDMNES toolbox", today=date.today(), calcroot=calcroot, fn_in=fn_in, fn_out=fn_out, fn_ext="txt", estart_sel="", estart="-20.", efermi_sel="", efermi="-5.36", spin="", core_sel="!", core="!", hole_sel="", hole="0.5", conv_const="!", conv_sel="", ecent="25.0", elarg="20.0", gamma_max="10.0", gamma_type="Gamma_fix", gauss_sel="", gaussian="0.9", ) else: self.opts = opts if calcroot is not None: self.opts["calcroot"] = calcroot self.opts["fn_in"] = "{}.{}".format(calcroot, self.opts["fn_ext"]) self.opts["fn_out"] = "{}_conv{}.{}".format( calcroot, self.opts["spin"], self.opts["fn_ext"] ) if fn_in is not None: self.opts["calcroot"] = fn_in[:-4] self.opts["fn_in"] = fn_in self.opts["fn_out"] = "{}_conv{}.{}".format( fn_in[:-4], self.opts["spin"], self.opts["fn_ext"] ) if fn_out is not None: self.opts["fn_out"] = fn_out # then check all options self.checkopts() def checkopts(self): if (self.opts["calcroot"] is None) or (self.opts["fn_in"] is None): raise NameError("missing 'calcroot' or 'fn_in'") if self.opts["estart"] == "!": self.opts["estart_sel"] = "!" if self.opts["efermi"] == "!": self.opts["efermi_sel"] = "!" if self.opts["spin"] == "up": self.opts["core_sel"] = "" self.opts["core"] = "2 !spin up" elif self.opts["spin"] == "down": self.opts["core_sel"] = "" self.opts["core"] = "1 !spin down" elif self.opts["spin"] == "": self.opts["core_sel"] = "!" elif self.opts["spin"] == "both": raise NameError('spin="both" not implemented!') else: self.opts["spin"] = "" self.opts["core_sel"] = "!" self.opts["core"] = "!" if self.opts["hole"] == "!": self.opts["hole_sel"] = "!" if self.opts["conv_const"] == "!": self.opts["conv_sel"] = "!" else: self.opts["conv_sel"] = "" if self.opts["gamma_type"] == "Gamma_fix": pass elif self.opts["gamma_type"] == "Gamma_var": pass else: raise NameError('gamma_type="Gamma_fix"/"Gamma_var"') if self.opts["gaussian"] == "!": self.opts["gauss_sel"] = "!" else: self.opts["gauss_sel"] = "" # update the output file name self.opts["fn_out"] = "{}_conv{}.{}".format( self.opts["calcroot"], self.opts["spin"], self.opts["fn_ext"] ) def setopt(self, opt, value): self.opts[opt] = value self.checkopts()
[docs] def wfdmfile(self): """ write a simple fdmfile.txt to enable the convolution first makes a copy of previous fdmfile.txt if not already done """ if os.path.exists("fdmfile.bak"): print("fdmfile.bak exists, good") else: subprocess.call("cp fdmfile.txt fdmfile.bak", shell=True) print("copied fdmfile.txt to fmdfile.bak") # s = Template( "!fdmfile.txt automatically created by ${creator} on ${today} (for convolution)\n\ !--------------------------------------------------------------------!\n\ ! Number of calculations\n\ 1\n\ ! FOR CONVOLUTION STEP\n\ convfile.txt\n\ !--------------------------------------------------------------------!\n\ " ) outstr = s.substitute(self.opts) f = open("fdmfile.txt", "w") f.write(outstr) f.close()
def wconvfile(self): s = Template( """ !FDMNES convolution file\n\ !created by ${creator} on ${today}\n\ ! Calculation\n\ ${fn_in}\n\ Conv_out\n\ ${fn_out}\n\ ${estart_sel}Estart\n\ ${estart_sel}${estart}\n\ ${efermi_sel}Efermi\n\ ${efermi_sel}${efermi}\n\ ${core_sel}Selec_core\n\ ${core_sel}${core}\n\ ${hole_sel}Gamma_hole\n\ ${hole_sel}${hole}\n\ ${conv_sel}Convolution\n\ ${conv_sel}${ecent} ${elarg} ${gamma_max} !Ecent Elarg Gamma_max\n\ ${conv_sel}${gamma_type}\n\ ${gauss_sel}Gaussian\n\ ${gauss_sel}${gaussian} !Gaussian conv for experimental res\n\ """ ) outstr = s.substitute(self.opts) f = open("convfile.txt", "w") f.write(outstr) f.close()
[docs] def run(self): """ runs fdmnes """ self.wfdmfile() # write fdmfile.txt self.wconvfile() # write convfile.txt try: subprocess.call("fdmnes", shell=True) except OSError: print("check 'fdmnes' executable exists!")
# LARCH PLUGIN # def registerLarchPlugin(): return (MODNAME, {"glinbroad": glinbroad}) if __name__ == "__main__": # tests/examples in xraysloth/examples/convolution1D_tests.py pass