Search code examples
pythonnumpymatplotlibnumpy-slicing

Slice list of 2D points for plotting with matplotlib


I've got a list of two-dimensional points represented in a numpy style array:

lines = np.array([
    [[1,1],  [2,3]],    # line 1 (x,y) -> (x,y)
    [[-1,1], [-2,2]],   # line 2 (x,y) -> (x,y)
    [[1,-1], [2,-7]]    # line 3 (x,y) -> (x,y)
])

I'd like to plot these lines with matplotlib in the simplest form possible. However, most of matplotlib's methods expect points to be represented component wise like ([x1, x2, x3, ...], [y1, y2, y3, ...]) instead of point wise.

I managed to get the slices right for a quiver type plot with vectors:

x = lines[:,0,0]
y = lines[:,0,1]
u = lines[:,1,0]
v = lines[:,1,1]

# quiver([X, Y], U, V, [C], **kwargs)
plt.quiver(x, y, u, v, color=['r','b','g'], scale=1, scale_units='xy', angles='xy')
plt.xticks(np.arange(-10, 10, 1))
plt.yticks(np.arange(-10, 10, 1))
plt.grid()
plt.show()

But it's my impression that quiver isn't really the right type of plot for my use case. Especially the scale, scale_units and angles arguments took me quite a while to figure out. Without them, vectors were displayed "wrong".

I'd rather use the simpler plot, in this signature

plot([x], y, [fmt], [x2], y2, [fmt2], ..., **kwargs)

which for each line to plot expects the two x coordinates followed by the y coordinates.

So for the lines above, we need

rearranged = [ [1,2],[1,3], [-1,-2],[1,2], [1,2],[-1,-7] ]
#              ^line1       ^line2         ^line3

If I add these to a call of plot statically, it's exactly what I need.

Question is, how can I slice or rearrange my initial lines array?

All I've got so far is

lines[:,[0,0],[0,1]]

but that only gives me the x coordinates of each line:

[[ 1  1]
 [-1  1]
 [ 1 -1]]

Solution

  • Use np.ndarray.T

    x, y = lines.T
    
    plt.plot(x, y)
    plt.show()
    

    Result:

    plot

    Transposed array:

    array([[[ 1, -1,  1],
            [ 2, -2,  2]],
    
           [[ 1,  1, -1],
            [ 3,  2, -7]]])