Pylab / Matplotlib图形

时间:2013-06-23 17:22:36

标签: python list matplotlib

我需要在Pylab中创建一个图形,我将沿Y轴绘制彩色点。 y轴从0到100。 我还有一个包含100个元素的列表,元素是+1或-1。 该列表必须与图表的Y轴相对应。

例如,如果列表中的第五个元素是+1,我需要在Y轴上的y = 5上绘制一个绿点。如果列表中的第五个元素是-1,则该点必须为红色。

我必须为列表中的所有元素执行此操作。

我在Pylab中绘制了简单的图形,但在这种情况下我完全迷失了。 任何帮助将不胜感激。谢谢!

1 个答案:

答案 0 :(得分:2)

import matplotlib.pyplot as plt
import numpy as np

data = np.array([1,1,-1,-1,1])
cmap = np.array([(1,0,0), (0,1,0)])
uniqdata, idx = np.unique(data, return_inverse=True)

N = len(data)
fig, ax = plt.subplots()
plt.scatter(np.zeros(N), np.arange(1, N+1), s=100, c=cmap[idx])
plt.grid()
plt.show()

产量

enter image description here


<强>解释

如果你打印出np.unique(data, return_inverse=True),你会看到它返回一个数组元组:

In [71]: np.unique(data, return_inverse=True)
Out[71]: (array([-1,  1]), array([1, 1, 0, 0, 1]))

第一个数组表示data中的唯一值是-1和1.第二个数组在data为-1的地方指定值0,在data为1的地方指定值1基本上,np.unique允许我们将[1,1,-1,-1,1]转换为[1, 1, 0, 0, 1]。现在cmap[idx]是一个RGB值数组:

In [74]: cmap[idx]
Out[74]: 
array([[0, 1, 0],
       [0, 1, 0],
       [1, 0, 0],
       [1, 0, 0],
       [0, 1, 0]])

这是所谓的"fancy indexing" on NumPy arrays的应用程序。 cmap[0]cmap的第一行。 cmap[1]cmap的第二行。 cmap[idx]是一个数组,cmap[idx]中的第i个元素是cmap[idx[i]]。因此,您最终将cmap[idx]作为2D数组,其中第i行为cmap[idx[i]]。因此cmap[idx]可以被认为是RGB颜色值的序列。


如果您有多个点并且希望在列中绘制它们,我能想到的最简单的方法是为每个ax.scatter列表调用data

import matplotlib.pyplot as plt
import numpy as np

def plot_data(ax, data, xval):
    N = len(data)
    uniqdata, idx = np.unique(data, return_inverse=True)
    ax.scatter(np.ones(N)*xval, np.arange(1, N+1), s=100, c=cmap[idx])

cmap = np.array([(1,0,0), (0,1,0)])
fig, ax = plt.subplots()

data = np.array([1,1,-1,-1,1])
data2 = np.array([1,-1,1,1,-1])

plot_data(ax, data, 0)
plot_data(ax, data2, 1)

plt.grid()
plt.show()

enter image description here

关于这一点的好处是它相对容易理解。关于这一点的坏处是它不止一次地调用ax.scatter。如果您有大量数据集,则整理数据和调用ax.scatter一次会更有效。 Matplotlib的速度更快,但代码更复杂:

import matplotlib.pyplot as plt
import numpy as np
import itertools as IT

def plot_dots(ax, datasets):
    N = sum(len(data) for data in datasets)
    x = np.fromiter(
        (i for i, data in enumerate(datasets) for j in np.arange(len(data))),
        dtype='float', count=N)
    y = np.fromiter(
        (j for data in datasets for j in np.arange(1, len(data)+1)),
        dtype='float', count=N)
    c = np.fromiter(
        (val for data in datasets
         for rgb in cmap[np.unique(data, return_inverse=True)[-1]]
         for val in rgb),
        dtype='float', count=3*N).reshape(-1,3)
    ax.scatter(x, y, s=100, c=c)

cmap = np.array([(1,0,0), (0,1,0)])
fig, ax = plt.subplots()

N = 100
datasets = [np.random.randint(2, size=5) for i in range(N)]

plot_dots(ax, datasets)

plt.grid()
plt.show()

enter image description here


参考文献: