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?

Accepted Answer

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
```

Plot 2-dimensional NumPy array using specific columns

Solution

Problem