Coverage for larch/xafs/rebin_xafs.py: 11%

76 statements  

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

1import numpy as np 

2 

3from larch import Group, Make_CallArgs, parse_group_args 

4from larch.math import (index_of, interp1d, 

5 remove_dups, remove_nans, remove_nans2) 

6from .xafsutils import ktoe, etok, TINY_ENERGY 

7 

8@Make_CallArgs(["energy", "mu"]) 

9def sort_xafs(energy, mu=None, group=None, fix_repeats=True, remove_nans=True, overwrite=True): 

10 """sort energy, mu pair of XAFS data so that energy is monotonically increasing 

11 

12 Arguments 

13 --------- 

14 energy input energy array 

15 mu input mu array 

16 group output group 

17 fix_repeats bool, whether to fix repeated energies [True] 

18 remove_nans bool, whether to fix nans [True] 

19 overwrite bool, whether to overwrite arrays [True] 

20 

21 Returns 

22 ------- 

23 None 

24 

25 if overwrite is False, a group named 'sorted' will be created 

26 in the output group, with sorted energy and mu arrays 

27 

28 (if the output group is None, _sys.xafsGroup will be written to) 

29 

30 """ 

31 energy, mu, group = parse_group_args(energy, members=('energy', 'mu'), 

32 defaults=(mu,), group=group, 

33 fcn_name='sort_xafs') 

34 

35 indices = np.argsort(energy) 

36 new_energy = energy[indices] 

37 new_mu = mu[indices] 

38 

39 if fix_repeats: 

40 new_energy = remove_dups(new_energy, tiny=TINY_ENERGY) 

41 if remove_nans: 

42 if (len(np.where(~np.isfinite(new_energy))[0]) > 0 or 

43 len(np.where(~np.isfinite(new_mu))[0] > 0)): 

44 new_energy, new_mu = remove_nans2(new_energy, new_mu) 

45 

46 if not overwrite: 

47 group.sorted = Group(energy=new_energy, mu=new_mu) 

48 else: 

49 group.energy = new_energy 

50 group.mu = new_mu 

51 return 

52 

53 

54@Make_CallArgs(["energy", "mu"]) 

55def rebin_xafs(energy, mu=None, group=None, e0=None, pre1=None, pre2=-30, 

56 pre_step=2, xanes_step=None, exafs1=15, exafs2=None, 

57 exafs_kstep=0.05, method='centroid'): 

58 """rebin XAFS energy and mu to a 'standard 3 region XAFS scan' 

59 

60 Arguments 

61 --------- 

62 energy input energy array 

63 mu input mu array 

64 group output group 

65 e0 energy reference -- all energy values are relative to this 

66 pre1 start of pre-edge region [1st energy point] 

67 pre2 end of pre-edge region, start of XANES region [-30] 

68 pre_step energy step for pre-edge region [2] 

69 xanes_step energy step for XANES region [see note] 

70 exafs1 end of XANES region, start of EXAFS region [15] 

71 exafs2 end of EXAFS region [last energy point] 

72 exafs_kstep k-step for EXAFS region [0.05] 

73 method one of 'boxcar', 'centroid' ['centroid'] 

74 

75 Returns 

76 ------- 

77 None 

78 

79 A group named 'rebinned' will be created in the output group, with the 

80 following attributes: 

81 energy new energy array 

82 mu mu for energy array 

83 e0 e0 copied from current group 

84 

85 (if the output group is None, _sys.xafsGroup will be written to) 

86 

87 Notes 

88 ------ 

89 1 If the first argument is a Group, it must contain 'energy' and 'mu'. 

90 See First Argrument Group in Documentation 

91 

92 2 If xanes_step is None, it will be found from the data as E0/25000, 

93 truncated down to the nearest 0.05: xanes_step = 0.05*max(1, int(e0/1250.0)) 

94 

95 

96 3 The EXAFS region will be spaced in k-space 

97 

98 4 The rebinned data is found by determining which segments of the 

99 input energy correspond to each bin in the new energy array. That 

100 is, each input energy is assigned to exactly one bin in the new 

101 array. For each new energy bin, the new value is selected from the 

102 data in the segment as either 

103 a) linear interpolation if there are fewer than 3 points in the segment. 

104 b) mean value ('boxcar') 

105 c) centroid ('centroid') 

106 

107 """ 

108 energy, mu, group = parse_group_args(energy, members=('energy', 'mu'), 

109 defaults=(mu,), group=group, 

110 fcn_name='rebin_xafs') 

111 

112 if e0 is None: 

113 e0 = getattr(group, 'e0', None) 

114 

115 if e0 is None: 

116 raise ValueError("need e0") 

117 

118 if pre1 is None: 

119 pre1 = pre_step*int((min(energy) - e0)/pre_step) 

120 

121 if exafs2 is None: 

122 exafs2 = max(energy) - e0 

123 

124 # determine xanes step size: 

125 # find mean of energy difference within 10 eV of E0 

126 nx1 = index_of(energy, e0-10) 

127 nx2 = index_of(energy, e0+10) 

128 de_mean = np.diff(energy[nx1:nx1]).mean() 

129 if xanes_step is None: 

130 xanes_step = 0.05 * max(1, int(e0 / 1250.0)) # E0/25000, round down to 0.05 

131 

132 # create new energy array from the 3 segments (pre, xanes, exafs) 

133 en = [] 

134 for start, stop, step, isk in ((pre1, pre2, pre_step, False), 

135 (pre2, exafs1, xanes_step, False), 

136 (exafs1, exafs2, exafs_kstep, True)): 

137 if isk: 

138 start = etok(start) 

139 stop = etok(stop) 

140 

141 npts = 1 + int(0.1 + abs(stop - start) / step) 

142 reg = np.linspace(start, stop, npts) 

143 if isk: 

144 reg = ktoe(reg) 

145 en.extend(e0 + reg[:-1]) 

146 

147 # find the segment boundaries of the old energy array 

148 bounds = [index_of(energy, e) for e in en] 

149 mu_out = [] 

150 err_out = [] 

151 

152 j0 = 0 

153 for i in range(len(en)): 

154 if i == len(en) - 1: 

155 j1 = len(energy) - 1 

156 else: 

157 j1 = int((bounds[i] + bounds[i+1] + 1)/2.0) 

158 if i == 0 and j0 == 0: 

159 j0 = index_of(energy, en[0]-5) 

160 # if not enough points in segment, do interpolation 

161 if (j1 - j0) < 3: 

162 jx = j1 + 1 

163 if (jx - j0) < 3: 

164 jx += 1 

165 

166 val = interp1d(energy[j0:jx], mu[j0:jx], en[i]) 

167 err = mu[j0:jx].std() 

168 if np.isnan(val): 

169 j0 = max(0, j0-1) 

170 jx = min(len(energy), jx+1) 

171 val = interp1d(energy[j0:jx], mu[j0:jx], en[i]) 

172 err = mu[j0:jx].std() 

173 else: 

174 if method.startswith('box'): 

175 val = mu[j0:j1].mean() 

176 else: 

177 val = (mu[j0:j1]*energy[j0:j1]).mean()/energy[j0:j1].mean() 

178 mu_out.append(val) 

179 err_out.append(mu[j0:j1].std()) 

180 j0 = j1 

181 

182 newname = group.__name__ + '_rebinned' 

183 group.rebinned = Group(energy=np.array(en), mu=np.array(mu_out), 

184 delta_mu=np.array(err_out), e0=e0, 

185 __name__=newname) 

186 return