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

1#!/usr/bin/env python 

2# -*- coding: utf-8 -*- 

3 

4""" 

5Generic discrete convolution 

6 

7Description 

8----------- 

9 

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. 

15 

16Resources 

17--------- 

18 

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> 

22 

23""" 

24 

25from __future__ import print_function 

26 

27__author__ = "Mauro Rovezzi" 

28__email__ = "mauro.rovezzi@gmail.com" 

29__credits__ = "Marius Retegan" 

30__version__ = "2023.5" 

31 

32DEBUG = 0 

33 

34import os, sys 

35import subprocess 

36from optparse import OptionParser 

37from datetime import date 

38from string import Template 

39import numpy as np 

40 

41from .lineshapes import gaussian, lorentzian 

42from .utils import polyfit 

43 

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 

61 

62 

63def lin_gamma(ene, gamma_hole=0.5, linbroad=None): 

64 """returns constant or linear energy-dependent broadening 

65 

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 

94 

95 

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)""" 

98 

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)) 

102 

103 # Set gamma to the gamma_hole below the cutoff energy. 

104 mask = np.where(e < e_cut) 

105 gammas[mask] = gamma_hole 

106 

107 return gammas 

108 

109 

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 

112 

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" 

129 

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)) 

134 

135 ids = x_ext > e_cut 

136 

137 x1 = x_ext[ids] 

138 x1[0] = e_cut 

139 

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]) 

143 

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 

152 

153 return y_conv 

154 

155 

156def conv(x, y, gammas, e_cut=None, kernel="gaussian"): 

157 """linear broadening 

158 

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" 

172 

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 

178 

179 if x.shape != gammas.shape: 

180 print("Error: 'gammas' array does not have the same shape of 'x'") 

181 return 0 

182 

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) 

188 

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)) 

192 

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 

221 

222 

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) 

226 

227 

228glinbroad.__doc__ = conv.__doc__ 

229 

230### CONVOLUTION WITH FDMNES VIA SYSTEM CALL ### 

231class FdmnesConv(object): 

232 """Performs convolution with FDMNES within Python""" 

233 

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() 

279 

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 ) 

321 

322 def setopt(self, opt, value): 

323 self.opts[opt] = value 

324 self.checkopts() 

325 

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() 

349 

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() 

379 

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!")