Coverage for pesummary/core/plots/figure.py: 77.8%
45 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-12-09 22:34 +0000
1# Licensed under an MIT style license -- see LICENSE.md
3from pesummary.utils.decorators import try_latex_plot
4from matplotlib.figure import Figure as MatplotlibFigure
6__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"]
9def figure(*args, gca=True, **kwargs):
10 """Extension of the matplotlib.pyplot.figure function
11 """
12 from matplotlib import pyplot
14 try:
15 _ = kwargs.pop("FigureClass", False)
16 except KeyError:
17 pass
19 kwargs["FigureClass"] = Figure
20 fig = pyplot.figure(*args, **kwargs)
21 if gca:
22 return fig, fig.gca()
23 return fig
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
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
46def _close(self):
47 """Extension of the matplotlib.pyplot.close function
48 """
49 from matplotlib.pyplot import close
51 for ax in self.axes[::-1]:
52 ax.set_xscale('linear')
53 ax.set_yscale('linear')
54 ax.cla()
55 close(self)
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
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)
72 def close(self):
73 """Close the plot
74 """
75 _close(self)
77 @try_latex_plot
78 def savefig(self, *args, **kwargs):
79 return super(Figure, self).savefig(*args, **kwargs)
81 @try_latex_plot
82 def tight_layout(self, *args, **kwargs):
83 return super(Figure, self).tight_layout(*args, **kwargs)