code.py 文件源码

def kmeans():

    T=20
    n_train = 500
    mix_weights = (0.2,0.5,0.3)

    #create a class object:
    km = kmeans_gaussian(n_train,mix_weights)

    # set k-values
    kvalues = range(2,6)
    colors = ['blue','green','red','black','yellow']

    #store the cluster assignments:
    k_backup = [3,5]
    cluster_assgn35 = []

    plt.figure()
    for i in range(len(kvalues)):
        km.set_k(kvalues[i])
        km.initialize_cluster_centers()
        km.train(T)

        plt.plot(range(1,T+1),km.objective,colors[i])

        #store cluster assignments for k=3,5
        if kvalues[i] in k_backup:
            cluster_assgn35.append(km.cluster_assgn[:,0])

    plt.xticks(range(1,T+1))
    plt.xlabel('Iterations')
    plt.ylabel('Objective')
    plt.title('Objective vs Iteration for K = [2,3,4,5]')
    plt.legend(['K = %d'%i for i in kvalues])
    # plt.savefig('hw4_1a_kmean_obj')
    # plt.show()

    #plot part b:
    for i in range(2):
        plt.figure()
        colors_arr = [colors[int(x)] for x in cluster_assgn35[i]]
        plt.scatter(km.data[:,0],km.data[:,1],c=colors_arr)
        plt.xlabel('Dimension 1')
        plt.ylabel('Dimension 2')
        plt.title('Scatter plot with cluster assignment for K=%d'%k_backup[i])
        plt.savefig('hw4_2_k%d.png'%k_backup[i])
        plt.show()

#############################
#### PART B - MATRIX FACT
#############################