matplotlib subplot2grid doesn't display correctly - python-2.7

Question

I'm using subplot2grid to display graphs. However, not all subplots are being displayed. Obviously it has to do with the if statement.
However, in my complete code I need those if statements because depending on some conditions plots have diffent formats. I want all 3 subplots to be displayed (one for each i). However, the first one is missing. How to display it correctly?
Here is the simplified code:
import matplotlib.pyplot as plt
fig=plt.figure()
for i in xrange(0,3):
if i==1:
ax=plt.subplot2grid((3,1),(i,0))
ax.plot([1,2],[1,2])
fig.autofmt_xdate()
else:
ax=plt.subplot2grid((3,1),(i,0), rowspan=2)
ax.plot([1,2],[1,2])
fig.autofmt_xdate()
plt.show()

Answer 1

I would just use the gridspec module from matplotlib. Then you can set the width/height ratios directly.
Then you can do something like this:
import numpy
from matplotlib import gridspec
import matplotlib.pyplot as plt
def do_plot_1(ax):
ax.plot([0.25, 0.5, 0.75], [0.25, 0.5, 0.75], 'k-')
def do_plot_2(ax):
ax.plot([0.25, 0.5, 0.75], [0.25, 0.5, 0.75], 'g--')
fig = plt.figure(figsize=(6, 4))
gs = gridspec.GridSpec(nrows=3, ncols=1, height_ratios=[2, 1, 2])
for n in range(3):
ax = fig.add_subplot(gs[n])
if n == 1:
do_plot_1(ax)
else:
do_plot_2(ax)
fig.tight_layout()
To use plt.subplot2grid, you'd need to effectively do something like this:
fig = plt.figure(figsize=(6, 4))
ax1 = plt.subplot2grid((5,1), (0, 0), rowspan=2)
ax2 = plt.subplot2grid((5,1), (2, 0), rowspan=1)
ax3 = plt.subplot2grid((5,1), (3, 0), rowspan=2)
Since you have two axes with a rowspan=2, your grid needs to be 2+1+2 = 5 blocks tall.