Coverage for tcvx21/plotting/plot_comparison_m.py: 96%

1""" 

2Routines for making plots to compare results for single observables 

3""" 

4import matplotlib.pyplot as plt 

5import tcvx21 

6plt.style.use(tcvx21.style_sheet) 

7 

8from .save_figure_m import savefig 

9from .plot_comparison_1d_m import plot_1D_comparison 

10from .plot_comparison_2d_m import plot_2D_comparison, plot_2D_comparison_simulation_only 

11 

12def plot_comparison(field_direction: str, diagnostic: str, observable: str, 

13 experimental_data: dict, simulation_data: dict, 

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

15 output_path = None, debug: bool = False, **kwargs): 

16 """ 

17 Plots a comparison of experimental_data against simulation_data for a single 

18 field direction, diagnostic and observable 

19 

20 simulation_data should be dictionary of the form 

21 {'simulation_name': {'forward_field': Record(), 'reversed_field': Record}} 

22 """ 

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

24 

25 if reference.is_empty: 

26 if debug: print(f"Missing experimental data for {field_direction}:{diagnostic}:{observable}") 

27 reference = list(simulation_data.values())[0][field_direction].get_observable(diagnostic, observable) 

28 simulation_only = True 

29 else: 

30 simulation_only = False 

31 

32 dimensionality = reference.dimensionality 

33 

34 keys = dict(field_direction=field_direction, diagnostic=diagnostic, observable=observable) 

35 kwargs = {**kwargs, **keys} 

36 

37 if dimensionality == 1: 

38 fig, ax = plot_1D_comparison(**kwargs, 

39 experimental_data=experimental_data, 

40 simulation_data=simulation_data) 

41 

42 else: 

43 if simulation_only: 

44 fig, ax = plot_2D_comparison_simulation_only(**kwargs, 

45 simulation_data=simulation_data) 

46 

47 else: 

48 fig, ax = plot_2D_comparison(**kwargs, 

49 experimental_data=experimental_data, 

50 simulation_data=simulation_data) 

51 

52 if output_path is None and save: 

53 output_path = tcvx21.results_dir / 'observables_fig' / diagnostic / \ 

54 f"{diagnostic}_{observable}_{field_direction}.png" 

55 

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

57 

58 return fig, ax