Coverage for /Users/Newville/Codes/xraylarch/larch/math/convolution1D.py: 15%
190 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-11-09 10:08 -0600
« prev ^ index » next coverage.py v7.3.2, created at 2023-11-09 10:08 -0600
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
4"""
5Generic discrete convolution
7Description
8-----------
10 This is a manual (not optimized!) implementation of discrete 1D
11 convolution intended for spectroscopy analysis. The difference with
12 commonly used methods is the possibility to adapt the convolution
13 kernel for each convolution point, e.g. change the FWHM of the
14 Gaussian kernel as a function of the energy scale.
16Resources
17---------
19.. [WPconv] <http://en.wikipedia.org/wiki/Convolution#Discrete_convolution>
20.. [Fisher] <http://homepages.inf.ed.ac.uk/rbf/HIPR2/convolve.htm>
21.. [GP1202] <http://glowingpython.blogspot.fr/2012/02/convolution-with-numpy.html>
23"""
25from __future__ import print_function
27__author__ = "Mauro Rovezzi"
28__email__ = "mauro.rovezzi@gmail.com"
29__credits__ = "Marius Retegan"
30__version__ = "2023.5"
32DEBUG = 0
34import os, sys
35import subprocess
36from optparse import OptionParser
37from datetime import date
38from string import Template
39import numpy as np
41from .lineshapes import gaussian, lorentzian
42from .utils import polyfit
44def get_ene_index(ene, cen, hwhm):
45 """returns the min/max indexes for array ene at (cen-hwhm) and (cen+hwhm)
46 very similar to index_of in larch
47 """
48 try:
49 if (cen - hwhm) <= min(ene):
50 ene_imin = 0
51 else:
52 ene_imin = max(np.where(ene < (cen - hwhm))[0])
53 if (cen + hwhm) >= max(ene):
54 ene_imax = len(e) - 1
55 else:
56 ene_imax = min(np.where(ene > (cen + hwhm))[0])
57 return ene_imin, ene_imax
58 except:
59 print("index not found for {0} +/- {1}".format(cen, hwhm))
60 return None, None
63def lin_gamma(ene, gamma_hole=0.5, linbroad=None):
64 """returns constant or linear energy-dependent broadening
66 Parameters
67 ----------
68 ene : energy array in eV
69 gamma_hole : initial full width at half maximum in eV
70 linbroad : list of 3-elements giving
71 'final full width at half maximum'
72 'start energy point of the linear increase'
73 'end energy point of the linear increase'
74 """
75 w = np.ones_like(ene)
76 if linbroad is None:
77 return w * gamma_hole
78 else:
79 try:
80 fwhm2 = linbroad[0]
81 e1 = linbroad[1]
82 e2 = linbroad[2]
83 except:
84 raise ValueError("wrong format for linbroad")
85 for en, ee in enumerate(ene):
86 if ee < e1:
87 w[en] *= gamma_hole
88 elif ee <= e2:
89 wlin = gamma_hole + (ee - e1) * (fwhm2 - gamma_hole) / (e2 - e1)
90 w[en] *= wlin
91 elif ee >= e2:
92 w[en] *= fwhm2
93 return w
96def atan_gamma(e, e_cut=0, e_cent=30, e_larg=30, gamma_hole=0.5, gamma_max=15):
97 """Arctangent energy dependent broadening as implemented in FDMNES (described in the "Convolution" section of the manual)"""
99 f = (e - e_cut) / e_cent
100 a = np.pi * gamma_max * (f - 1 / f**2) / (3 * e_larg)
101 gammas = gamma_hole + gamma_max * (0.5 + 1 / np.pi * np.arctan(a))
103 # Set gamma to the gamma_hole below the cutoff energy.
104 mask = np.where(e < e_cut)
105 gammas[mask] = gamma_hole
107 return gammas
110def conv_fast(x, y, gammas, e_cut=None, kernel="gaussian", num=501, step=0.1):
111 """A significantly faster version of the `conv` function
113 Parameters
114 ----------
115 x : x-axis (energy)
116 y : f(x) to convolve with g(x) kernel, y(energy)
117 kernel : convolution kernel, g(x)
118 'gaussian'
119 'lorentzian'
120 gammas : the full width half maximum in eV for the kernel
121 broadening. It is an array of size 'e' with constants or
122 an energy-dependent values determined by a function as
123 'lin_gamma()' or 'atan_gamma()'
124 """
125 assert e_cut is not None, "starting energy for the convolution not given"
126 assert (
127 "lor" or "gauss" in kernel.lower()
128 ), "the kernel should be either Lorentzian or Gaussian"
130 # Extend the X-axis array.
131 start = x[-1]
132 stop = start + num * step
133 x_ext = np.append(x, np.arange(start + step, stop, step))
135 ids = x_ext > e_cut
137 x1 = x_ext[ids]
138 x1[0] = e_cut
140 # Extend the intensity array by coping the last value `num - 1` times.
141 y = y[x > e_cut]
142 y = np.append(y, np.ones(num - 1) * y[-1])
144 y_conv = np.zeros_like(x)
145 for i, (xi, gamma) in enumerate(zip(x, gammas)):
146 gamma = gamma / 2.0
147 if "gauss" in kernel.lower():
148 ky = gaussian(x1, center=xi, sigma=gamma)
149 elif "lor" in kernel.lower():
150 ky = lorentzian(x1, center=xi, sigma=gamma)
151 y_conv[i] = np.sum(ky * y) / np.pi
153 return y_conv
156def conv(x, y, gammas, e_cut=None, kernel="gaussian"):
157 """linear broadening
159 Parameters
160 ----------
161 x : x-axis (energy)
162 y : f(x) to convolve with g(x) kernel, y(energy)
163 kernel : convolution kernel, g(x)
164 'gaussian'
165 'lorentzian'
166 gammas : the full width half maximum in eV for the kernel
167 broadening. It is an array of size 'e' with constants or
168 an energy-dependent values determined by a function as
169 'lin_gamma()' or 'atan_gamma()'
170 """
171 assert e_cut is not None, "starting energy for the convolution not given"
173 f = y[:]*1.0
174 z = np.zeros_like(f)
175 # ief = index_nearest(x, e_cut)
176 ief = np.argmin(np.abs(x - e_cut))
177 f[0:ief] *= 0
179 if x.shape != gammas.shape:
180 print("Error: 'gammas' array does not have the same shape of 'x'")
181 return 0
183 # linear fit upper part of the spectrum to avoid border effects
184 # polyfit => pf
185 lpf = int(len(x) / 2.0)
186 cpf = polyfit(x[-lpf:], f[-lpf:], 1, reverse=False)
187 fpf = np.polynomial.Polynomial(cpf)
189 # extend upper energy border to 3*fhwm_e[-1]
190 xstep = x[-1] - x[-2]
191 xup = np.append(x, np.arange(x[-1] + xstep, x[-1] + 3 * gammas[-1], xstep))
193 for n in range(len(f)):
194 # from now on I change e with eup
195 eimin, eimax = get_ene_index(xup, xup[n], 1.5 * gammas[n])
196 if (eimin is None) or (eimax is None):
197 if DEBUG:
198 raise IndexError("e[{0}]".format(n))
199 if len(range(eimin, eimax)) % 2 == 0:
200 kx = xup[eimin : eimax + 1] # odd range centered at the convolution point
201 else:
202 kx = xup[eimin:eimax]
203 ### kernel ###
204 hwhm = gammas[n] / 2.0
205 if "gauss" in kernel.lower():
206 ky = gaussian(kx, center=xup[n], sigma=hwhm)
207 elif "lor" in kernel.lower():
208 ky = lorentzian(kx, center=xup[n], sigma=hwhm)
209 else:
210 raise ValueError("convolution kernel '{0}' not implemented".format(kernel))
211 ky = ky / ky.sum() # normalize
212 zn = 0
213 lk = len(kx)
214 for mf, mg in zip(range(-int(lk / 2), int(lk / 2) + 1), range(lk)):
215 if ((n + mf) >= 0) and ((n + mf) < len(f)):
216 zn += f[n + mf] * ky[mg]
217 elif (n + mf) >= 0:
218 zn += fpf(xup[n + mf]) * ky[mg]
219 z[n] = zn
220 return z
223def glinbroad(e, mu, gammas=None, e_cut=None):
224 """gaussian linear convolution in Larch"""
225 return conv(e, mu, gammas=gammas, kernel="gaussian", e_cut=e_cut)
228glinbroad.__doc__ = conv.__doc__
230### CONVOLUTION WITH FDMNES VIA SYSTEM CALL ###
231class FdmnesConv(object):
232 """Performs convolution with FDMNES within Python"""
234 def __init__(self, opts=None, calcroot=None, fn_in=None, fn_out=None):
235 if opts is None:
236 self.opts = dict(
237 creator="FDMNES toolbox",
238 today=date.today(),
239 calcroot=calcroot,
240 fn_in=fn_in,
241 fn_out=fn_out,
242 fn_ext="txt",
243 estart_sel="",
244 estart="-20.",
245 efermi_sel="",
246 efermi="-5.36",
247 spin="",
248 core_sel="!",
249 core="!",
250 hole_sel="",
251 hole="0.5",
252 conv_const="!",
253 conv_sel="",
254 ecent="25.0",
255 elarg="20.0",
256 gamma_max="10.0",
257 gamma_type="Gamma_fix",
258 gauss_sel="",
259 gaussian="0.9",
260 )
261 else:
262 self.opts = opts
263 if calcroot is not None:
264 self.opts["calcroot"] = calcroot
265 self.opts["fn_in"] = "{}.{}".format(calcroot, self.opts["fn_ext"])
266 self.opts["fn_out"] = "{}_conv{}.{}".format(
267 calcroot, self.opts["spin"], self.opts["fn_ext"]
268 )
269 if fn_in is not None:
270 self.opts["calcroot"] = fn_in[:-4]
271 self.opts["fn_in"] = fn_in
272 self.opts["fn_out"] = "{}_conv{}.{}".format(
273 fn_in[:-4], self.opts["spin"], self.opts["fn_ext"]
274 )
275 if fn_out is not None:
276 self.opts["fn_out"] = fn_out
277 # then check all options
278 self.checkopts()
280 def checkopts(self):
281 if (self.opts["calcroot"] is None) or (self.opts["fn_in"] is None):
282 raise NameError("missing 'calcroot' or 'fn_in'")
283 if self.opts["estart"] == "!":
284 self.opts["estart_sel"] = "!"
285 if self.opts["efermi"] == "!":
286 self.opts["efermi_sel"] = "!"
287 if self.opts["spin"] == "up":
288 self.opts["core_sel"] = ""
289 self.opts["core"] = "2 !spin up"
290 elif self.opts["spin"] == "down":
291 self.opts["core_sel"] = ""
292 self.opts["core"] = "1 !spin down"
293 elif self.opts["spin"] == "":
294 self.opts["core_sel"] = "!"
295 elif self.opts["spin"] == "both":
296 raise NameError('spin="both" not implemented!')
297 else:
298 self.opts["spin"] = ""
299 self.opts["core_sel"] = "!"
300 self.opts["core"] = "!"
301 if self.opts["hole"] == "!":
302 self.opts["hole_sel"] = "!"
303 if self.opts["conv_const"] == "!":
304 self.opts["conv_sel"] = "!"
305 else:
306 self.opts["conv_sel"] = ""
307 if self.opts["gamma_type"] == "Gamma_fix":
308 pass
309 elif self.opts["gamma_type"] == "Gamma_var":
310 pass
311 else:
312 raise NameError('gamma_type="Gamma_fix"/"Gamma_var"')
313 if self.opts["gaussian"] == "!":
314 self.opts["gauss_sel"] = "!"
315 else:
316 self.opts["gauss_sel"] = ""
317 # update the output file name
318 self.opts["fn_out"] = "{}_conv{}.{}".format(
319 self.opts["calcroot"], self.opts["spin"], self.opts["fn_ext"]
320 )
322 def setopt(self, opt, value):
323 self.opts[opt] = value
324 self.checkopts()
326 def wfdmfile(self):
327 """write a simple fdmfile.txt to enable the convolution
328 first makes a copy of previous fdmfile.txt if not already done"""
329 if os.path.exists("fdmfile.bak"):
330 print("fdmfile.bak exists, good")
331 else:
332 subprocess.call("cp fdmfile.txt fdmfile.bak", shell=True)
333 print("copied fdmfile.txt to fmdfile.bak")
334 #
335 s = Template(
336 "!fdmfile.txt automatically created by ${creator} on ${today} (for convolution)\n\
337!--------------------------------------------------------------------!\n\
338! Number of calculations\n\
339 1\n\
340! FOR CONVOLUTION STEP\n\
341convfile.txt\n\
342!--------------------------------------------------------------------!\n\
343"
344 )
345 outstr = s.substitute(self.opts)
346 f = open("fdmfile.txt", "w")
347 f.write(outstr)
348 f.close()
350 def wconvfile(self):
351 s = Template(
352 """
353!FDMNES convolution file\n\
354!created by ${creator} on ${today}\n\
355!
356Calculation\n\
357${fn_in}\n\
358Conv_out\n\
359${fn_out}\n\
360${estart_sel}Estart\n\
361${estart_sel}${estart}\n\
362${efermi_sel}Efermi\n\
363${efermi_sel}${efermi}\n\
364${core_sel}Selec_core\n\
365${core_sel}${core}\n\
366${hole_sel}Gamma_hole\n\
367${hole_sel}${hole}\n\
368${conv_sel}Convolution\n\
369${conv_sel}${ecent} ${elarg} ${gamma_max} !Ecent Elarg Gamma_max\n\
370${conv_sel}${gamma_type}\n\
371${gauss_sel}Gaussian\n\
372${gauss_sel}${gaussian} !Gaussian conv for experimental res\n\
373"""
374 )
375 outstr = s.substitute(self.opts)
376 f = open("convfile.txt", "w")
377 f.write(outstr)
378 f.close()
380 def run(self):
381 """runs fdmnes"""
382 self.wfdmfile() # write fdmfile.txt
383 self.wconvfile() # write convfile.txt
384 try:
385 subprocess.call("fdmnes", shell=True)
386 except OSError:
387 print("check 'fdmnes' executable exists!")