Coverage for tcvx21/analysis/fit_tables_m.py: 93%

1""" 

2Makes tables for fitted profiles 

3""" 

4import numpy as np 

5import matplotlib.pyplot as plt 

6from pathlib import Path 

7import pandas as pd 

8from collections import namedtuple 

9import tcvx21 

10from tcvx21.analysis import fit_exponential_decay, fit_eich_profile 

11from tcvx21.analysis.fit_eich_profile_m import fit_eich_profile, eich_profile 

12from tcvx21.plotting.labels_m import label_subplots 

13 

14from tcvx21.plotting.labels_m import add_twinx_label, make_labels, add_x_zero_line, add_y_zero_line, format_yaxis 

15from tcvx21.quant_validation.ricci_metric_m import calculate_normalised_ricci_distance_from_observables 

16from tcvx21.plotting import savefig 

17 

18 

19def decay_rate_row(experimental_data, simulation_data, field_direction, diagnostic, observable, row_label, fit_range=(0, 1.0), plot=False): 

20 """ 

21 Make rows of the decay rate table (and plot if requested) 

22 """ 

23 Row = namedtuple('Row', ('TCV', *simulation_data.keys())) 

24 results = [] 

25 

26 fd = '+' if field_direction == 'forward_field' else '-' 

27 if plot: plt.figure() 

28 

29 ref = experimental_data[field_direction].get_observable(diagnostic, observable) 

30 _, (K, K_err) = fit_exponential_decay(ref.positions, ref.values, ref.errors, fit_range, plot=f'TCV{fd}' if plot else '') 

31 results.append(f"{K.to('cm').magnitude:.1f}$\\pm${K_err.to('cm').magnitude:.1f}" if K is not None else '-') 

32 

33 for key, case in simulation_data.items(): 

34 sim = case[field_direction].get_observable(diagnostic, observable) 

35 if sim.is_empty: 

36 K = None 

37 else: 

38 _, (K, K_err) = fit_exponential_decay(sim.positions, sim.values, sim.errors if sim.has_errors else None, 

39 fit_range, plot=f'{key}{fd}' if plot else '') 

40 results.append(f"{K.to('cm').magnitude:.1f}$\\pm${K_err.to('cm').magnitude:.1f}" if K is not None else '-') 

41 

42 if plot: 

43 plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left') 

44 plt.title(f"{row_label}$_{{{fd}}}$") 

45 

46 return f"$\\lambda_{{{row_label}}}$$^{{{fd}}}$", Row(*results) 

47 

48 

49def make_decay_rate_table(experimental_data, simulation_data, diagnostics, observables, labels, fit_range=(0, 1.0), plot=False, 

50 output_path: Path = None): 

51 """ 

52 Makes a table giving the exponential decay rates near the separatrix for a given list of diagnostics and observables 

53 Writes a LaTeX table as output 

54 """ 

55 assert np.array(diagnostics).size == np.array(observables).size == np.array(labels).size 

56 

57 rows, extended_labels = [], [] 

58 for field_direction in ['forward_field', 'reversed_field']: 

59 for diagnostic, observable, label in zip(diagnostics, observables, labels): 

60 

61 label, row = decay_rate_row(experimental_data, simulation_data, 

62 field_direction, diagnostic, observable, label, fit_range=fit_range, plot=plot) 

63 

64 rows.append(row) 

65 extended_labels.append(label) 

66 

67 df = pd.DataFrame(rows, index=extended_labels) 

68 if output_path is None: 

69 output_path = tcvx21.results_dir/'tables'/'exp_decay.tex' 

70 df.to_latex(output_path, float_format="{:0.1f}".format, escape=False, column_format=f"l{(len(simulation_data)+1)*'c'}") 

71 

72 return df 

73 

74def style_eich_plot(ax, reference): 

75 add_x_zero_line(ax) 

76 add_y_zero_line(ax) 

77 ax.set_xlim(reference.xlim) 

78 reference.apply_plot_limits(ax) 

79 ax.set_xticks([-1, 0, 1, 2]) 

80 plt.legend() 

81 

82def append_to_list_as_str(list_in, values=None, units=None): 

83 if values is None: 

84 list_in.append("-") 

85 else: 

86 values = values.to(units).magnitude 

87 list_in.append(f"{values[0]:.1f}$\\pm${values[1]:.1f}") 

88 

89def eich_comparison(experimental_data, simulation_data, field_direction, diagnostic, observable, ax): 

90 lambda_q_list, S_list = [], [] 

91 

92 reference = experimental_data[field_direction].get_observable(diagnostic, observable) 

93 reference.plot(ax) 

94 

95 peak_heat_flux, lambda_q, spreading, background_heat_flux, shift = fit_eich_profile(reference.positions, reference.values) 

96 append_to_list_as_str(lambda_q_list, lambda_q, 'mm') 

97 append_to_list_as_str(S_list, spreading, 'mm') 

98 

99 ax.plot(reference.positions, eich_profile(reference.positions, 

100 peak_heat_flux[0], lambda_q[0], spreading[0], background_heat_flux[0], shift[0]), 

101 linestyle='--', linewidth=0.75*plt.rcParams['lines.linewidth']) 

102 

103 field_direction_string, diagnostic_string, observable_string = make_labels(field_direction, diagnostic, reference) 

104 add_twinx_label(ax, field_direction_string) 

105 ax.set_title(f"{diagnostic_string}: {observable_string}") 

106 

107 for key, case in simulation_data.items(): 

108 

109 simulation = case[field_direction].get_observable(diagnostic, observable) 

110 if simulation.is_empty: 

111 append_to_list_as_str(lambda_q_list, None) 

112 append_to_list_as_str(S_list, None) 

113 continue 

114 line = simulation.plot(ax, trim_to_x=reference.xlim, linestyle='-') 

115 

116 peak_heat_flux, lambda_q, spreading, background_heat_flux, shift = fit_eich_profile(simulation.positions, simulation.values) 

117 append_to_list_as_str(lambda_q_list, lambda_q, 'mm') 

118 append_to_list_as_str(S_list, spreading, 'mm') 

119 

120 ax.plot(simulation.positions, eich_profile(simulation.positions, 

121 peak_heat_flux[0], lambda_q[0], spreading[0], background_heat_flux[0], shift[0]), 

122 linestyle='--', linewidth=0.75*plt.rcParams['lines.linewidth'], 

123 color=simulation.color) 

124 

125 d = calculate_normalised_ricci_distance_from_observables(simulation=simulation, experiment=reference) 

126 line.set_label(f"{line.get_label()}: d={d:3.2f}") 

127 

128 style_eich_plot(ax, reference) 

129 

130 return lambda_q_list, S_list 

131 

132def make_eich_fit_comparison(experimental_data, simulation_data, diagnostic='LFS-IR', observable='q_parallel', 

133 table_output_path: Path = None, fig_output_path: Path = None): 

134 

135 Row = namedtuple('Row', ('TCV', *simulation_data.keys())) 

136 fig, axs = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(3.25, 5)) 

137 

138 rows, labels = [], [] 

139 for field_direction, ax in zip(['forward_field', 'reversed_field'], axs.flatten()): 

140 fd = '+' if field_direction == 'forward_field' else '-' 

141 labels.append(f'$\\lambda_{{q}}^{{{fd}}}$') 

142 labels.append(f'$S^{{{fd}}}$') 

143 

144 (lambda_q_list, S_list) = eich_comparison(experimental_data, simulation_data, 

145 field_direction, diagnostic, observable, ax) 

146 

147 rows.append(Row(*lambda_q_list)) 

148 rows.append(Row(*S_list)) 

149 

150 axs[0].set_xlabel('') 

151 axs[1].set_title('') 

152 axs[0].set_ylim(bottom=0) 

153 axs[1].set_ylim(bottom=0) 

154 label_subplots(axs) 

155 

156 plt.draw() 

157 for ax in axs.flatten(): 

158 format_yaxis(ax) 

159 

160 if fig_output_path is None: 

161 fig_output_path = tcvx21.results_dir/'analysis_fig'/'eich_fit.png' 

162 savefig(fig, fig_output_path) 

163 

164 df = pd.DataFrame(rows, index=labels) 

165 if table_output_path is None: 

166 table_output_path = tcvx21.results_dir/'tables'/'lambda_q.tex' 

167 df.to_latex(table_output_path, float_format="{:0.1f}".format, escape=False, column_format=f"l{(len(simulation_data)+1)*'c'}") 

168 

169 return df