Matplotlib draw connections between two sets of 2D points efficiently

plot_connected_points.py

from matplotlib.collections import LineCollection

def plot_connected_pairs(a,b,*args,**kwargs):
    """
    Draw lines between two arrays of 2D points.
    a and b must be the same shape, both [N,2] arrays to be plotted.
    
    Parameters:
    -----------
    a:  [N,2] array of points to plot from
    b:  [N,2] array of points to plot to
    
    Any other arguments or keyword arguments are passed 
    to LineCollection directly.
    
    Returns:
    --------
    LineCollection of connecting lines
    
    """
    segs = np.einsum('nij -> inj', np.stack([a,b]))    
    line_segments = LineCollection(segs, *args, **kwargs)
    plt.gca().add_collection(line_segments)
    return line_segments