5 Best Ways to Return a Figure in Python Using Matplotlib

by

💡 Problem Formulation: In data visualization tasks, often there is a need to create functions that not only generate graphics using matplotlib, a popular Python plotting library, but also return these figures so they can be further modified or saved externally. This article introduces various methods to return a matplotlib figure from a function. The user calls the function, which generates a figure and returns it, ready for display or export (e.g., fig = my_plotting_function()).

Method 1: Using the Figure Object

Matplotlib’s Figure object is the core element to plot graphs and can be generated inside a function and returned for external use. This approach enables the caller to manipulate the figure once it’s returned.

Here’s an example:

import matplotlib.pyplot as plt

def create_figure():
    fig, ax = plt.subplots() 
    ax.plot([0, 1, 2], [0, 1, 0])
    return fig

fig = create_figure()
fig.show()

Output: A line plot with points (0,0), (1,1), (2,0).

This example demonstrates creating a basic line plot within a function and returning the Figure object. You can display the plot with fig.show() or save it with fig.savefig('filename.png').

Method 2: Returning an Axes object

Instead of returning the entire figure, you can return the Axes object for more fine-tuned control over the individual plot elements.

Here’s an example:

import matplotlib.pyplot as plt

def create_axes():
    fig, ax = plt.subplots()
    ax.plot([3, 2, 1], [1, 2, 3])
    return ax

ax = create_axes()
plt.show()

Output: A line plot with points (3,1), (2,2), (1,3).

This method returns an Axes object which can be manipulated post-return, giving the user the ability to adjust the axes, labels, title, etc., before displaying the plot with plt.show().

Method 3: Saving the Figure to a Buffer

Sometimes you might want to return a figure as an image file. In this case, you can save the matplotlib figure to a buffer using BytesIO.

Here’s an example:

import matplotlib.pyplot as plt
from io import BytesIO

def save_figure_to_buffer():
    fig, ax = plt.subplots()
    ax.plot([1, 2, 3], [3, 2, 1])
    buffer = BytesIO()
    fig.savefig(buffer, format='png')
    buffer.seek(0)
    return buffer

buffer = save_figure_to_buffer()
# Now you can write the buffer content to a file or manipulate it further.

Output: Buffer object contains a PNG image of the plot.

In this code snippet, the figure is saved to an in-memory buffer as a PNG image, which can be useful for sending the image over a network or saving it without a need for a temporary file.

Method 4: Embedding Figures in applications

When integrating a matplotlib figure into a larger application such as a web frontend or a GUI, you can return a reference to the figure which can then be embedded directly into various user interface toolkits.

Here’s an example:

import matplotlib.pyplot as plt

def embed_figure():
    fig, ax = plt.subplots()
    ax.plot([-1, -2, -3], [-3, -2, -1])
    return fig.canvas

canvas = embed_figure()
# canvas can now be used to integrate the plot into toolkits like Tkinter, WxPython, etc.

Output: The canvas associated with the figure.

This approach provides the canvas object tied to the figure, which can then be used to display the plot inside various frameworks and toolkits, enhancing the integration of matplotlib figures within custom applications.

Bonus One-Liner Method 5: Returning a Figure with a Lambda Function

A concise way to create and return figures in Python is by using lambda functions, enabling the creation of an ad-hoc function without defining a formal function.

Here’s an example:

import matplotlib.pyplot as plt

create_and_show_figure = lambda: plt.figure().add_subplot().plot([4, 5, 6])

create_and_show_figure()
plt.show()

Output: A plot with default settings.

This one-liner lambda function will create a new figure and plot the provided data points straight away when called, ideal for quick and temporary plotting tasks.

Summary/Discussion

  • Method 1: Using the Figure Object. Provides the full feature-set of matplotlib’s plotting capabilities. Strengths: High level of customization and control. Weaknesses: Requires more lines of code.
  • Method 2: Returning an Axes object. Offers a balance between control over the plotting area and simplicity. Strengths: Easier manipulation of individual plots within a figure. Weaknesses: Less intuitive for managing multiple subplots.
  • Method 3: Saving the Figure to a Buffer. Useful for programmatically handling figure image data. Strengths: Great for saving or transmitting data without a filesystem. Weaknesses: Additional complexity when processing the buffer.
  • Method 4: Embedding Figures in applications. Ideal for integrating matplotlib within broader applications. Strengths: Seamless integration into various GUI frameworks. Weaknesses: Requires knowledge of the target UI toolkit.
  • Method 5: Returning a Figure with a Lambda Function. Fastest way to plot when working interactively. Strengths: Very concise code. Weaknesses: Limited customization and not suitable for complex plotting scenarios.