How are existing subplot axes retrieved using subplot2grid/gridspec in Matplotlib?

Question

I have problems accessing existing subplot in a Matplotlib figure when specifying plot locations using gridspec directly, or alternatively, subplot2grid. Regular subplot specs, e.g. add_subplot(211), returns existing axes if any. Using gridspec/subplot2grid seems to destroy any existing axes. How do I retrieve the existing axes objects using gridspec/subplot2grid? Is this intended behaviour or am I missing something here? I would like a solution where I do not have to define own placeholders for the axes objects.

Example:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

x = np.linspace(0,10,100)
y1 = np.cos(x)
y2 = np.sin(x)

fig = plt.figure()
ax = fig.add_subplot(211)
ax.plot(x,y1, '-b')
ax = fig.add_subplot(212)
ax.plot(x,y2, '-b')
ax = fig.add_subplot(211) #here, the existing axes object is retrieved
ax.plot(x,y2, '-r')

fig = plt.figure()
gs = gridspec.GridSpec(2,1)
ax = fig.add_subplot(gs[0,0])
ax.plot(x,y1, '-b')
ax = fig.add_subplot(gs[1,0])
ax.plot(x,y2, '-b')
# using gridspec (or subplot2grid), existing axes
# object is apparently deleted
ax = fig.add_subplot(gs[0,0])
ax.plot(x,y2, '-r')

plt.show()

Answer 1

This is actually a subtle bug with the magic of how add_subplot determines if a an axes exists. It boils down to this fact:

In [220]: gs[0, 0] == gs[0, 0]
Out[220]: False

which is because gridspec.__getitem__ returns a new object every time you call it and SubplotSpec does not overload __eq__ so python checks 'is this the same object in memory' when searching for existing axes.

That is what is wrong, however my naive attempt to fix it by adding a __eq__ to a SubplotSpec and monkey patching matplotlib.gridspec.SubplotSpec didn't work (I don't have time to figure out why), but if you add

def __eq__(self, other):
    return all((self._gridspec == other._gridspec,
                self.num1 == other.num1,
                self.num2 == other.num2))

to class SubplotSpec(object): in matplotlib/gridspec.py ~L380 and re-installing from source works as expected.

PR to fix this which seems to break all sorts of other things.