Coverage for tcvx21/analysis/fit_eich_profile_m.py: 92%

1""" 

2Routines to fit Eich-type heat flux profiles 

3""" 

4import numpy as np 

5import matplotlib.pyplot as plt 

6import scipy.special 

7from scipy.optimize import curve_fit 

8from tcvx21 import Quantity 

9import warnings 

10 

11 

12def erfc(x): 

13 """ 

14 Complimentary error function, with handling of Quantity objects 

15 Only Dimensionless quantities can be handled 

16 erfc(x) = 1 - erf(x) 

17 """ 

18 if isinstance(x, Quantity): 

19 return scipy.special.erfc(x.to('').magnitude) 

20 else: 

21 return scipy.special.erfc(x) 

22 

23 

24def eich_profile(position, peak_heat_flux, lambda_q, spreading_factor, background_heat_flux, shift): 

25 """ 

26 Fits an Eich profile to heat flux data, using upstream-mapped 

27 distance to remove the explicit flux-expansion factor 

28 

29 See Eich Nuclear Fusion 2013. Using the OMP-mapped distance as position 

30 means that you should set 'f_x' = 1 

31 

32 You can pass the parameters to the profile either all as dimensional 

33 Quantity objects, or all as dimensionless raw arrays/floats 

34 

35 position: OMP-mapped radial distance [mm] 

36 peak_heat_flux: peak heat flux [MW/m^2] 

37 lambda_q: heat flux decay width [mm] 

38 spreading_factor: spreading factor [mm] 

39 background_heat_flux: background heat-flux [MW/m^2] 

40 """ 

41 with warnings.catch_warnings(): 

42 warnings.simplefilter("ignore", category=RuntimeWarning) 

43 return peak_heat_flux / 2.0 \ 

44 * np.exp((spreading_factor / (2.0 * lambda_q)) ** 2 - (position - shift) / lambda_q) \ 

45 * erfc(spreading_factor / (2.0 * lambda_q) - (position - shift) / spreading_factor) \ 

46 + background_heat_flux 

47 

48 

49def fit_eich_profile(position, heat_flux, errors=None, plot=False): 

50 """ 

51 Uses non-linear least-squares to fit an Eich profile to heat flux data 

52 

53 Should pass 

54 position: OMP-mapped radial distance [mm] 

55 heat_flux: total heat flux [MW/m^2] 

56 

57 We use the Levenberg-Marquardt algorithm (default) for curve fitting 

58 

59 Note that the algorithm takes into account the error if provided 

60 """ 

61 

62 if errors is None: 

63 errors = np.ones_like(position) 

64 has_errors = False 

65 else: 

66 errors = errors.to('MW/m^2').magnitude 

67 has_errors = True 

68 

69 heat_flux = heat_flux.to('MW/m^2').magnitude 

70 

71 mask = np.logical_and.reduce((~np.isnan(heat_flux), ~np.isnan(errors), errors/heat_flux > 1E-2)) 

72 

73 position = position.to('mm').magnitude[mask] 

74 heat_flux = heat_flux[mask] 

75 errors = errors[mask] 

76 

77 popt, pcov = curve_fit(eich_profile, 

78 xdata=position, 

79 ydata=heat_flux, 

80 sigma=errors, 

81 absolute_sigma=has_errors, 

82 p0=(np.max(heat_flux), 5, 2, np.mean((heat_flux[0], heat_flux[-1])), 0.0), 

83 method='lm', 

84 xtol=1E-6 

85 ) 

86 

87 peak_heat_flux, lambda_q, spreading_factor, background_heat_flux, shift = popt 

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

89 peak_heat_flux_err, lambda_q_err, spreading_err, background_heat_flux_err, shift_err = perr 

90 

91 if plot: 

92 plt.plot(Quantity(position, 'mm'), Quantity(heat_flux, 'MW/m^2'), label='signal') 

93 if has_errors: 

94 plt.fill_between(Quantity(position, 'mm'), Quantity(heat_flux + errors, 'MW/m^2'), Quantity(heat_flux - errors, 'MW/m^2')) 

95 plt.plot(Quantity(position, 'mm'), Quantity(eich_profile(position, *popt), 'MW/m^2'), label='eich') 

96 plt.legend() 

97 plt.title("Eich profile fitting") 

98 

99 return Quantity((peak_heat_flux, peak_heat_flux_err), 'MW/m^2'), \ 

100 Quantity((lambda_q, lambda_q_err), 'mm'), \ 

101 Quantity((spreading_factor, spreading_err), 'mm'), \ 

102 Quantity((background_heat_flux, background_heat_flux_err), 'MW/m^2'), \ 

103 Quantity((shift, shift_err), 'mm')