Coverage for tcvx21/analysis/fit_exp_decay_m.py: 95%

1""" 

2Routines to fit exponential decay functions to observables 

3""" 

4import numpy as np 

5import matplotlib.pyplot as plt 

6from scipy.optimize import curve_fit 

7from tcvx21 import Quantity 

8 

9 

10def exponential_decay(x, amplitude, decay_rate): 

11 """A curve given simple exponential decay""" 

12 return amplitude * np.exp(-x/decay_rate) 

13 

14 

15def linear_fit(x, intercept, gradient): 

16 return gradient * x + intercept 

17 

18# K, A_log = np.polyfit(position, np.log(values), 1) 

19# A = np.exp(A_log) 

20# return A, K 

21 

22 

23def fit_exponential_decay(position, values, errors=None, fit_range=(-np.inf, np.inf), plot='', fit_logarithm=False): 

24 """ 

25 Fit an exponential decay rate to some observable 

26 If you supply plot as a non-empty string, then this will be used to label a fit of the plot 

27 fit_range should be given as (min, max), in cm 

28 

29 If supplying errors, use these in an absolute sense (i.e. don't rescale the errors to calculate a relative weighting 

30 of the points -- use the errors as directly setting the uncertainty of the fit) 

31 """ 

32 

33 position = position.to('cm').magnitude 

34 

35 if errors is None: 

36 errors = np.ones_like(position) 

37 has_errors = False 

38 elif fit_logarithm: 

39 raise NotImplementedError("No implementation for fit_logarithm with errors") 

40 else: 

41 errors = errors.to(values.units).magnitude 

42 has_errors = True 

43 

44 mask = np.logical_and.reduce((position > fit_range[0], position < fit_range[1], 

45 ~np.isnan(values), ~np.isnan(errors))) 

46 

47 units = values.units 

48 values = values.magnitude[mask] 

49 position = position[mask] 

50 errors = errors[mask] 

51 

52 if len(position) < 3: 

53 return (None, None), (None, None) 

54 

55 if fit_logarithm: 

56 popt, pcov = curve_fit(linear_fit, xdata=position, ydata=np.log(values)) 

57 perr = np.sqrt(np.diag(pcov)) 

58 

59 amplitude = np.exp(popt[0]) 

60 amplitude_error = amplitude * perr[0] 

61 decay_rate = -1/popt[1] 

62 decay_rate_error = np.abs(decay_rate * (perr[1]/popt[1])) 

63 

64 else: 

65 popt, pcov = curve_fit(exponential_decay, xdata=position, ydata=values, sigma=errors, 

66 absolute_sigma = has_errors, 

67 p0=([np.max(values), 1.0])) 

68 perr = np.sqrt(np.diag(pcov)) 

69 amplitude, decay_rate = popt 

70 amplitude_error, decay_rate_error = perr 

71 

72 if plot: 

73 plt.plot(position, exponential_decay(position, amplitude, decay_rate)) 

74 if has_errors: 

75 plt.errorbar(position, values, errors, label=plot) 

76 else: 

77 plt.plot(position, values, label=plot) 

78 plt.gca().set_yscale("log", nonpositive='clip') 

79 

80 return Quantity((amplitude, amplitude_error), units), Quantity((decay_rate, decay_rate_error), 'cm')