Coverage for larch/xafs/cauchy_wavelet.py: 95%

42 statements  

« prev     ^ index     » next       coverage.py v7.6.0, created at 2024-10-16 21:04 +0000

1#!/usr/bin/env python 

2# 

3# Cauchy Wavelet for EXAFS, adopted from 

4# matlab code from Munoz, Argoul, and Farges: 

5# 

6# CONTINUOUS CAUCHY WAVELET TRANSFORM OF EXAFS SIGNAL 

7# code freely downloaded from http://www.univ-mlv.fr/~farges/waw 

8# (c) 2000, Univ. Marne la Vallee, France 

9# 

10# please cite us of this code with: 

11# Munoz M., Argoul P. and Farges F. 

12# Continuous Cauchy wavelet transform analyses of 

13# EXAFS spectra: a qualitative approach. 

14# American Mineralogist 88, pp. 694-700 (2003). 

15# 

16# version history: 

17# 1999 Hans-Argoul : core wavelet algorithm 

18# 1999-2002 Argoul-Munoz : EXAFS adapation 

19# 2002 Farges : graphical and user interface 

20# 2003 Munoz : CPU optimizations and graphical updates 

21# 2003 Farges-Munoz : various fixes and web version 

22# 

23# 2014-Apr M Newville : translated to Python for Larch 

24 

25import numpy as np 

26from larch import Make_CallArgs, parse_group_args 

27from larch.math import complex_phase 

28from .xafsutils import set_xafsGroup 

29 

30@Make_CallArgs(["k" ,"chi"]) 

31def cauchy_wavelet(k, chi=None, group=None, kweight=0, rmax_out=10, 

32 nfft=2048, _larch=None): 

33 """ 

34 Cauchy Wavelet Transform for XAFS, following work of Munoz, Argoul, and Farges 

35 

36 Parameters: 

37 ----------- 

38 k: 1-d array of photo-electron wavenumber in Ang^-1 or group 

39 chi: 1-d array of chi 

40 group: output Group 

41 rmax_out: highest R for output data (10 Ang) 

42 kweight: exponent for weighting spectra by k**kweight 

43 nfft: value to use for N_fft (2048). 

44 

45 Returns: 

46 --------- 

47 None -- outputs are written to supplied group. 

48 

49 Notes: 

50 ------- 

51 Arrays written to output group: 

52 wcauchy complex cauchy wavelet(k, R) 

53 wcauchy_r uniform array of R, out to rmax_out. 

54 wcauchy_mag magnitude of wavelet(k, R) 

55 wcauchy_re real part of wavelet(k, R) 

56 wcauchy_im imaginary part of wavelet(k, R) 

57 

58 Supports First Argument Group convention (with group 

59 member names 'k' and 'chi') 

60 

61 """ 

62 k, chi, group = parse_group_args(k, members=('k', 'chi'), 

63 defaults=(chi,), group=group, 

64 fcn_name='cauchy_wavelet') 

65 

66 kstep = np.round(1000.*(k[1]-k[0]))/1000.0 

67 rstep = (np.pi/2048)/kstep 

68 rmin = 1.e-7 

69 rmax = rmax_out 

70 nrpts = int(np.round((rmax-rmin)/rstep)) 

71 nkout = len(k) 

72 if kweight != 0: 

73 chi = chi * k**kweight 

74 

75 # extend EXAFS to 1024 data points... 

76 NFT = int(nfft/2) 

77 if len(k) < NFT: 

78 knew = np.arange(NFT) * kstep 

79 xnew = np.zeros(NFT) * kstep 

80 xnew[:len(k)] = chi 

81 else: 

82 knew = k[:NFT] 

83 xnew = chi[:NFT] 

84 

85 # FT parameters 

86 freq = (1.0/kstep)*np.arange(nfft)/(2*nfft) 

87 omega = 2*np.pi*freq 

88 

89 # simple FT calculation 

90 tff = np.fft.fft(xnew, n= 2*nfft) 

91 

92 # scale parameter 

93 r = np.linspace(0, rmax, nrpts) 

94 r[0] = 1.e-19 

95 a = nrpts/(2*r) 

96 

97 # Characteristic values for Cauchy wavelet: 

98 cauchy_sum = np.log(2*np.pi) - np.log(1.0+np.arange(nrpts)).sum() 

99 

100 # Main calculation: 

101 out = np.zeros(nkout*nrpts, 

102 dtype='complex128').reshape(nrpts, nkout) 

103 for i in range(nrpts): 

104 aom = a[i]*omega 

105 aom[np.where(aom==0)] = 1.e-19 

106 filt = cauchy_sum + nrpts*np.log(aom) - aom 

107 tmp = np.conj(np.exp(filt))*tff[:nfft] 

108 out[i, :] = np.fft.ifft(tmp, 2*nfft)[:nkout] 

109 

110 group = set_xafsGroup(group, _larch=_larch) 

111 group.wcauchy_r = r 

112 group.wcauchy = out 

113 group.wcauchy_mag = np.sqrt(out.real**2 + out.imag**2) 

114 group.wcauchy_re = out.real 

115 group.wcauchy_im = out.imag