Coverage for pesummary/core/plots/figure.py: 77.8%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3from pesummary.utils.decorators import try_latex_plot 

4from matplotlib.figure import Figure as MatplotlibFigure 

5 

6__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

7 

8 

9def figure(*args, gca=True, **kwargs): 

10 """Extension of the matplotlib.pyplot.figure function 

11 """ 

12 from matplotlib import pyplot 

13 

14 try: 

15 _ = kwargs.pop("FigureClass", False) 

16 except KeyError: 

17 pass 

18 

19 kwargs["FigureClass"] = Figure 

20 fig = pyplot.figure(*args, **kwargs) 

21 if gca: 

22 return fig, fig.gca() 

23 return fig 

24 

25 

26def subplots( 

27 nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True, subplot_kw=None, 

28 gridspec_kw=None, **fig_kw 

29): 

30 """Extension of the matplotlib.pyplot.subplots class 

31 """ 

32 try: 

33 _ = fig_kw.pop("gca", False) 

34 except KeyError: 

35 pass 

36 

37 fig_kw["gca"] = False 

38 fig = figure(**fig_kw) 

39 axs = fig.subplots( 

40 nrows=nrows, ncols=ncols, sharex=sharex, sharey=sharey, squeeze=squeeze, 

41 subplot_kw=subplot_kw, gridspec_kw=gridspec_kw 

42 ) 

43 return fig, axs 

44 

45 

46def _close(self): 

47 """Extension of the matplotlib.pyplot.close function 

48 """ 

49 from matplotlib.pyplot import close 

50 

51 for ax in self.axes[::-1]: 

52 ax.set_xscale('linear') 

53 ax.set_yscale('linear') 

54 ax.cla() 

55 close(self) 

56 

57 

58class ExistingFigure(object): 

59 """An extension of the core matplotlib `~matplotlib.figure.Figure` 

60 """ 

61 def __new__(self, fig): 

62 fig.close = lambda: _close(fig) 

63 return fig 

64 

65 

66class Figure(MatplotlibFigure): 

67 """An extension of the core matplotlib `~matplotlib.figure.Figure` 

68 """ 

69 def __init__(self, *args, **kwargs): 

70 super(Figure, self).__init__(*args, **kwargs) 

71 

72 def close(self): 

73 """Close the plot 

74 """ 

75 _close(self) 

76 

77 @try_latex_plot 

78 def savefig(self, *args, **kwargs): 

79 return super(Figure, self).savefig(*args, **kwargs) 

80 

81 @try_latex_plot 

82 def tight_layout(self, *args, **kwargs): 

83 return super(Figure, self).tight_layout(*args, **kwargs)