Matplotlib slider loses interactivity upon imshow

Question

I've been trying to use Matplotlib's imshow function to display a heatmap which has an interactive slider below it to be eventually used for displaying specific data.

For the code itself, I have two renditions that both have a heatmap and an interactive slider along with a close button. The one not using imshow() works fine under small loads while still being able to use the slider. The code sample provided, using imshow(), works extremely well for large quantities of data but loses the sliders interactivity, though the button isn't affected?

Any help is appreciated!

import numpy as np
import matplotlib

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

from matplotlib.widgets import Button, Slider

def heatmap(data, ax=None,
            cbar_kw={}, cbarlabel="", cbar=True, **kwargs):
    """
    Create a heatmap from a numpy array and two lists of labels.

    Parameters
    ----------
    data
        A 2D numpy array of shape (N, M).
    ax
        A `matplotlib.axes.Axes` instance to which the heatmap is plotted.  If
        not provided, use current axes or create a new one.  Optional.
    cbar_kw
        A dictionary with arguments to `matplotlib.Figure.colorbar`.  Optional.
    cbarlabel
        The label for the colorbar.  Optional.
    cbar
        Draws the colorbar. Optional.
    **kwargs
        All other arguments are forwarded to `imshow`.
    """

    if not ax:
        ax = plt.gca()

    # Plot the heatmap
    im = ax.imshow(data, **kwargs)

    if cbar:
        # Create colorbar
        cbar = ax.figure.colorbar(im, ax=ax, **cbar_kw)
        cbar.ax.set_ylabel(cbarlabel, rotation=-90, va="bottom")

    # We want to show all ticks...
    ax.set_xticks(np.arange(data.shape[1]))
    ax.set_yticks(np.arange(data.shape[0]))

    # Let the horizontal axes labeling appear on top.
    ax.tick_params(top=False, bottom=True,
                   labeltop=False, labelbottom=True)

    # Turn spines off and create white grid.
    for edge, spine in ax.spines.items():
        spine.set_visible(False)

    '''
    ax.grid(which="minor", color="w", linestyle='-', linewidth=3)
    plt.setp(ax.get_xticklabels(), visible=False)
    plt.setp(ax.get_yticklabels(), visible=False)
    ax.tick_params(axis='both', which='both', length=0)
    '''
    # Adjust tick boundaries based on min/max
    ax.xaxis.set_ticks(np.arange(0, 1000, 200))
    ax.xaxis.set_major_formatter(ticker.FormatStrFormatter('%0.1f'))
    ax.yaxis.set_ticks(np.arange(0, 220, 20))
    ax.yaxis.set_major_formatter(ticker.FormatStrFormatter('%0.1f'))

    ax.grid(which="minor", color="w", linestyle='-', linewidth=3)

    return im, cbar


fig, ax = plt.subplots()

im, cb = heatmap(np.random.rand(220, 1000), ax=ax,
        cmap="jet", aspect='auto')

closeax = plt.axes([0.8, 0.025, 0.1, 0.04])
drpax = plt.axes([0.15, 0.025, 0.50, 0.03])

axcolor = 'lightgoldenrodyellow'

# Button location
button = Button(closeax, 'Close', color=axcolor, hovercolor='0.975')
def close_fig(evt):
    plt.close()
button.on_clicked(close_fig)

sdrp = Slider(drpax, 'Slider', 1, 10000, valinit=10000, valstep=1)

plt.ion()
plt.show()


while True:
    ax.cla()
    plt.cla()

    im, cb = heatmap(np.random.rand(220, 1000), ax=ax,
        cmap="jet", cbar=False, aspect='auto')

    fig.canvas.draw()
    plt.pause(0.001)

matplotlib.axes.Axes.imshow
matplotlib.pyplot.imshow
matplotlib.figure.Figure.colorbar
matplotlib.pyplot.colorbar