gmm.py 文件源码

def draw(X, pred, means, covariances, output):
    xp = cupy.get_array_module(X)
    for i in six.moves.range(2):
        labels = X[pred == i]
        if xp is cupy:
            labels = labels.get()
        plt.scatter(labels[:, 0], labels[:, 1], c=np.random.rand(3))
    if xp is cupy:
        means = means.get()
        covariances = covariances.get()
    plt.scatter(means[:, 0], means[:, 1], s=120, marker='s', facecolors='y',
                edgecolors='k')
    x = np.linspace(-5, 5, 1000)
    y = np.linspace(-5, 5, 1000)
    X, Y = np.meshgrid(x, y)
    for i in six.moves.range(2):
        Z = mlab.bivariate_normal(X, Y, np.sqrt(covariances[i][0]),
                                  np.sqrt(covariances[i][1]),
                                  means[i][0], means[i][1])
        plt.contour(X, Y, Z)
    plt.savefig(output)