How can I programmatically select a specific subplot in Matplotlib?

So in a figure where three vertical subplots have been added with add_subplot, how can I select let's say the middle one?

Right now I do this list comprehension:

[r[0] for r in sorted([[ax, ax.get_geometry()[2]] for ax in self.figure.get_axes()], key=itemgetter(1))]

where I can simply select the index I want, with the corresponding axes. Is there a more straightforward way of doing this?

Solution

From the matplotlib documentation:

If the figure already has a subplot with key (args, kwargs) then it will simply make that subplot current and return it.

Here's an example:

import matplotlib.pyplot as plt

fig = plt.figure()  
for vplot in [1,2,3]:
    ax = fig.add_subplot(3,1,vplot)
    ax.plot(range(10),range(10))

ax_again = fig.add_subplot(3,1,2)
ax_again.annotate("The middle one",xy=(7,5),xytext=(7,5))

plt.show()

The middle plot is called again so that it can be annotated.

What if I set the background with my original call, do I need to set it again when I get the subplot the second time?

Yes. The arguments and keywords for the original call are used to make a unique identifier. So for the figure to generate this unique identifier again, you need to pass the same arguments (grid definition, position) and keywords again. For example:

import matplotlib.pyplot as plt

fig = plt.figure()  
ax = fig.add_subplot(2,1,1,axisbg='red')
ax.plot(range(10),range(10))
ax = fig.add_subplot(2,1,2)
ax.plot(range(10),range(10))

ax_again = fig.add_subplot(2,1,1,axisbg='red')
ax_again.annotate("The top one",xy=(7,5),xytext=(7,5))

plt.show()

What if I use ax_again.change_geometry() ?

You would think change_geometry, e.g. from a 312 to a 422, would change how you use add_subplot, but it doesn't. There appears to be a bug or undefined behavior when you call change_geometry. The unique key that was original generated using the arguments and keywords, to the first add_subplot call, does not get updated. Therefore, if you want to get an axis back with an add_subplot call, you need to call add_subplot with the original arguments and keywords. For more info, follow this issue report: https://github.com/matplotlib/matplotlib/issues/429

My guess for now is that if you change any property of the subplot after generating it with add_subplot call, the unique will not be adjusted. So just use the original arguments and keywords, and hopefully this will work out.