Convert a matplotlib plot to its most minimal form (no text, labels, titles, legends)

minimal_plot.py

"""
********************************************************************************
Matplotlib can create some great plots, although it is often a hassle to 
customize said plots after their creation, particularly if you're creating a 
large number of them, or regenerating them for minor changes (e.g. scaling)
is inconvenient.

I've been experimenting with generating "minimal" plots, which contain only the
actual representation of the data.
Insofar as I can add titles/labels in LaTeX, this allows me to embed the plots
in PDFs without having to worry if resizing them will render parts unreadable, 
or introduce weird artifacts.
The usefulness of the following script is a bit niche, but I expect I'll refer
to it in the future.

How it works:
    0. You have plotted a graph, perhaps labelled it, added a legend etc.
    1. Search over the children of the main object (the figure)
    2. Hide objects that you don't want to be visible
    3. Save the minimal figure
    
Somewhat to my surprise, this can be accomplished pretty easily-- I didn't
find problems with weird hidden state or circular dependencies, at least
when trying this with my own graphs.
"""
import matplotlib as mpl
import matplotlib.pyplot as plt

# ******************************************************************************
# Creating an example (non-minimal) graph and saving it.
# ******************************************************************************
# Example data
xdata = [i for i in range(10)] 
y1 = [x**2 for x in xdata]
y2 = [x**2 - x for x in xdata]

# Creating the original plot
fig, ax = plt.subplots(figsize=(4,3))

# Plotting with label
line_1 = ax.plot(xdata, y1, label='first line')
line_2 = ax.plot(xdata, y2, label='second line')

# We add a legend to the graph; the way I do it is overkill here, but whatever
# Lines to show on legend and their labels
lines = [line_1, line_2]
labels = [i.get_label() for i in lines]

# Create the legend (that shows up on the main plot)
# These values are changeable, should create a legend outside to the
# right of the graph.
legend = fig.legend(
    lines, 
    labels, 
    loc='upper right',
    bbox_to_anchor=(0.98, 0.9 , 0.32, -0.102),
    mode='expand',
    ncol=1,
    bbox_transform=fig.transFigure,
)

# Formatting
ax.set_title('Lorem vs. Ipsum')
ax.set_xlabel('x-axis')
ax.set_yabel('not x-axis')

# Save the figure, ensuring that the legend doesn't get cut off
# using `bbox_extra_artists`
figpath = 'graph.png'
fig.savefig(figpath, bbox_extra_artists=[legend], bbox_inches='tight')

# ******************************************************************************
# Now, to undo all of that labelling nonsense
# ******************************************************************************
# The types of elements we will set to invisible
TO_HIDE = set([
    mpl.text.Text, 
    mpl.legend.Legend,
])

# Get children of the figure
children = fig.get_children()

# Build a list of child elements via BFS
found = []
ids  = set()
while children:
    child = children.pop()
    found.append(child)
    
    # I'm a bit paranoid about possible circular dependencies
    child_id = id(child)
    if child_id not in ids:
        ids.add(child_id)
        # Add child objects to queue
        children.extend(child.get_children())
    
# Hiding child elements according to some criteria (modify as desired)
for child in found:
    if type(child) in TO_HIDE:
        child.set_visible(False)

# We save the minimalist version under a different name
figpath = 'graph-minimal.png'
print(f"Saving to: {figpath}")
fig.savefig(figpath, bbox_extra_artists=[], bbox_inches='tight')