Coverage for gpaw/utilities/extrapolate.py: 8%

1import sys 

2import numpy as np 

3from ase.parallel import paropen 

4 

5 

6def extrapolate(x, y, n=-1.5, plot=0, reg=0, txt=None): 

7 """Extrapolation tool. Mainly intended for RPA correlation energies, 

8 but could be useful for other purposes. Fits a straight line to an 

9 expression of the form: y=b + alpha*x**n and extrapolates the result 

10 to infinite x. reg=N gives linear regression using the last N points in 

11 x. reg should be larger than 2""" 

12 

13 if txt is None: 

14 f = sys.stdout 

15 elif isinstance(txt, str): 

16 f = paropen(txt, 'a') 

17 else: 

18 f = txt 

19 assert len(x) == len(y) 

20 ext = [] 

21 print('Two-point extrapolation:', file=f) 

22 for i in range(len(x) - 1): 

23 alpha = (y[i] - y[i + 1]) / (x[i]**n - x[i + 1]**n) 

24 ext.append(y[i + 1] - alpha * x[i + 1]**n) 

25 print(' ', x[i], '-', x[i + 1], ':', ext[-1], file=f) 

26 print(file=f) 

27 

28 if plot: 

29 import matplotlib.pyplot as plt 

30 plt.plot(x**n, y, 'o-', label='Data') 

31 plt.xticks(x**n, [int(e) for e in x]) 

32 plt.axis([0, None, None, None]) 

33 if reg > 2: 

34 a = x[-reg:]**n 

35 b = y[-reg:] 

36 N = reg 

37 delta = N * np.sum(a**2) - (np.sum(a))**2 

38 A = (np.sum(a**2) * np.sum(b) - np.sum(a) * np.sum(a * b)) / delta 

39 B = (N * np.sum(a * b) - np.sum(a) * np.sum(b)) / delta 

40 sigma_y = (1 / (N - 2) * np.sum((b - A - B * a)**2))**0.5 

41 sigma_A = sigma_y * (np.sum(a**2) / delta)**0.5 

42 

43 print('Linear regression using last %s points:' % N, file=f) 

44 print(' Extrapolated result:', A, file=f) 

45 print(' Uncertainty:', sigma_A, file=f) 

46 print(file=f) 

47 if plot: 

48 print([a[0], 0], [A + B * a[0], A]) 

49 plt.plot([a[0], 0], [A + B * a[0], A], '--', label='Regression') 

50 plt.legend(loc='upper left') 

51 else: 

52 A = 0 

53 B = 0 

54 sigma_A = 0 

55 if plot: 

56 plt.show() 

57 plt.plot(x[1:], ext, 'o-', label='Two-point extrapolation') 

58 if reg > 2: 

59 plt.plot([x[-reg], x[-1]], [A, A], '--', label='Regression') 

60 plt.errorbar(x[-2], A, yerr=sigma_A, 

61 elinewidth=2.0, capsize=5, c='g') 

62 

63 plt.legend(loc='lower right') 

64 plt.show() 

65 if txt is not None: 

66 f.close() 

67 return ext, A, B, sigma_A