Plot 2-dimensional NumPy array using specific columns

Question

I have a 2D numpy array that's created like this:

data = np.empty((number_of_elements, 7))

Each row with 7 (or whatever) floats represents an object's properties. The first two for example are the x and y position of the object, the others are various properties that could even be used to apply color information to the plot.

I want to do a scatter plot from data, so that if p = data[i], an object is plotted as a point with p[:2] as its 2D position and with say p[2:4] as a color information (the length of that vector should determine a color for the point). Other columns should not matter to the plot at all.

How should I go about this?

unutbu · Accepted Answer · 2016-01-23 01:10:54Z

29

Setting up a basic matplotlib figure is easy:

import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)

Picking off the columns for x, y and color might look something like this:

N = 100
data = np.random.random((N, 7))
x = data[:,0]
y = data[:,1]
points = data[:,2:4]
# color is the length of each vector in `points`
color = np.sqrt((points**2).sum(axis = 1))/np.sqrt(2.0)
rgb = plt.get_cmap('jet')(color)

The last line retrieves the jet colormap and maps each of the float values (between 0 and 1) in the array color to a 3-tuple RGB value. There is a list of colormaps to choose from here. There is also a way to define custom colormaps.

Making a scatter plot is now straight-forward:

ax.scatter(x, y, color = rgb)
plt.show()
# plt.savefig('/tmp/out.png')    # to save the figure to a file

enter image description here

edited Jan 23, 2016 at 1:10

answered Nov 29, 2012 at 20:59

unutbu

886k197 gold badges1.9k silver badges1.7k bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

EnricoGiampieri Over a year ago

Complete as usual. I would just suggest to not use the jet colormap but a more continuous one, like summer or winter, which give a more clear sense of the transition, that can be distorted by the jet color spectrum.

Daniel · Accepted Answer · 2012-11-29 20:47:26Z

11

Not sure exactly what you are looking for in the plot, but you can slice 2D arrays like this:

>>> a
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> a[:,1]
array([1, 4, 7])
>>> a[:,1:3]
array([[1, 2],
       [4, 5],
       [7, 8]])

Then some matplot to take care of the plotting. If you find what you are looking for at the Matplotlib Gallery I can help you more.

answered Nov 29, 2012 at 20:47

Daniel

19.6k7 gold badges64 silver badges74 bronze badges

1 Comment

tamacun Over a year ago

Wow, that was so easy it's embarassing.. I need to look into slicing stuff.

Collectives™ on Stack Overflow

Plot 2-dimensional NumPy array using specific columns

2 Answers 2

1 Comment

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related