Coverage for tcvx21/plotting/tile2d_m.py: 91%

1import matplotlib.pyplot as plt 

2import tcvx21 

3plt.style.use(tcvx21.style_sheet) 

4 

5from pathlib import Path 

6from tcvx21 import Quantity 

7from tcvx21.quant_validation.ricci_metric_m import calculate_normalised_ricci_distance_from_observables 

8from tcvx21.plotting.labels_m import add_twinx_label, add_x_zero_line, make_observable_string, label_subplots 

9import numpy as np 

10from .save_figure_m import savefig 

11from typing import Union 

12from .tile2d_single_observable import plot_2D_observables 

13 

14 

15def tile2d(experimental_data, simulation_data, 

16 diagnostics: Union[np.ndarray, str, tuple], 

17 observables: Union[np.ndarray, str, tuple], 

18 labels: list, 

19 fig_width=7.5, fig_height_per_row=2.0, title_height=1.0, manual_title: str = '', 

20 extra_args=None, x_title=0.875, subplots_kwargs={}, offsets: np.ndarray = None, 

21 show: bool = False, save: bool = True, close: bool = True, 

22 output_path: Path = None, **kwargs): 

23 """ 

24 Plots multiple 2D observables 

25 extra_args should be a tuple or list of the length of the diagnostics and observables, which contains 

26 keyword arguments passed to each key-set 

27 """ 

28 assert len(diagnostics) == len(observables) 

29 assert len(labels) == len(observables) 

30 if extra_args is not None: 

31 assert len(extra_args) == len(observables) 

32 

33 nrows = 2 * len(observables) 

34 

35 fig, axs = plt.subplots(ncols=len(simulation_data) + 1, 

36 nrows=nrows, 

37 figsize=(fig_width, nrows * fig_height_per_row), 

38 sharex='col', sharey='row') 

39 

40 if subplots_kwargs: 

41 plt.subplots_adjust(**subplots_kwargs) 

42 

43 if offsets is None: 

44 offsets = len(observables) * [None] 

45 

46 for index, key in enumerate(zip(diagnostics, observables)): 

47 

48 if extra_args is not None: 

49 key_kwargs = extra_args[index] 

50 else: 

51 key_kwargs = {} 

52 

53 subaxs = axs[2 * index:2 * index + 2] 

54 label_subplots(subaxs, prefix=f"{index+1}") 

55 

56 _, cbar = plot_2D_observables(axs=subaxs, key=key, 

57 experimental_data=experimental_data, simulation_data=simulation_data, 

58 offset=offsets[index], 

59 cbar_pad=0.025, add_labels=False, **key_kwargs, **kwargs) 

60 

61 reference = experimental_data['forward_field'].get_observable(*key) 

62 compact_units = reference.compact_units 

63 units_string = f"{Quantity(1, compact_units).units:~P}" if compact_units else "-" 

64 if offsets[index] is not None: 

65 units_string = f"$10^{{{np.log10(offsets[index]):n}}}${units_string.lstrip('1')}" 

66 

67 fig.text(x_title, 

68 (subaxs[-1, 0].get_position().y1 + subaxs[-1, 1].get_position().y1) / 2, 

69 f"{labels[index]} [{units_string}]", va='center', rotation=270) 

70 

71 for ax, title in zip(axs[0], ['TCV', *simulation_data.keys()]): 

72 ax.set_title(title) 

73 

74 fig.text(0.5, 0.09, '$R^u - R^u_{sep}$ [cm]', ha='center') 

75 fig.text(0.02, 0.5, '$Z - Z_X$ [m]', va='center', rotation='vertical') 

76 

77 if output_path is None and save: 

78 

79 filename = '+'.join( 

80 [f'{diagnostic}_{observable}' for diagnostic, observable in zip(diagnostics, observables)]) 

81 

82 output_path = tcvx21.results_dir / 'summary_fig' / f"{filename.replace(' ', '_')}.png" 

83 

84 savefig(fig, output_path=output_path, show=show, close=close) 

85 

86 return fig, axs