quiver.py

import matplotlib.pyplot as plt
import io

def quiver_overlay(img, x, y, u, v, rescale, **kwargs):
    """Plots a quiver and return a superposable visualisation
    
    Parameters
    ----------
    img: ndarray,
        the input (2D) image array
    x: ndarray,
        x origins of the arrows
    y: ndarray,
        y origins of the arrows
    u: ndarray,
        x component of the arrows
    v: ndarray,
        y component of the arrows
    rescale: int,
        factor to rescale the overlay by with respect to the
        original image
    **kwargs are passed to the matplotib `quiver` method 
    
    Returns
    -------
    quivers : 2D ndarray with the same shape as img
    
    
    
    See Also
    --------
    
    matplotlib.pyplot.quiver 
    
    
    
    """
    
    ny, nx = img.shape
    dpi = max(nx, ny)

    fig = plt.Figure(figsize=(nx / dpi, ny / dpi), dpi=dpi)
    ax = fig.add_axes((0, 0, 1, 1))

    imax = ax.imshow(
        np.zeros((nx, ny)),
        interpolation='none',
        cmap="gray_r",
        extent=(0, nx, ny, 0)
    )
    ax.quiver(x, y, u, v, **kwargs)
    ax.set_axis_off()

    with io.BytesIO() as io_buf:
        fig.savefig(io_buf, format='raw', dpi=overscale*dpi)
        io_buf.seek(0)

        quivers = np.reshape(
            np.frombuffer(io_buf.getvalue(), dtype=np.uint8),
            newshape=(int(fig.bbox.bounds[3]),
                      int(fig.bbox.bounds[2]),
                      -1)
        )[..., 0] # keep only red channel
    plt.close(fig)
    
    return quivers


time = 10
quivers = []

dpi = 300
overscale = 1

for time in range(rsize.shape[0]) :
    
    img = rsize[time]
    x = x_grid
    y = y_grid
    u_ = u[time]
    v_ = -v[time]
    quivers.append(
        quiver_overlay(img, x, y, u_, v_, overscale, facecolor="k", headwidth=2)
    )