Plotting clusterings with matplotlib and numpy

plot_clusterings.py

# Plotting clusterings with matplotlib

import numpy as np
import matplotlib.pyplot as plt

def plot_clusterings(data, clusterings, k):
    """
    Args:
        data - n x d numpy array
        clusterings - 1 x n numpy array, values from 0 to k-1
        k - number of clusters
    """
    colormap = plt.cm.rainbow
    for c in range(k):
        color = colormap(float(c)/float(k))
        plt.scatter(data[clusterings==c, 0], data[clusterings==c, 1], c = color)
    plt.show()