kmeans/kmeans_example.py

Create a set and points and cluster it, printing the resulting cluster centers

## \example kmeans/kmeans_example.py
# Create a set and points and cluster it, printing the resulting cluster
# centers

from __future__ import print_function
import IMP
import IMP.kmeans
import sys

IMP.setup_from_argv(sys.argv, "kmeans example")


def add_points(km):
    # add data
    km.add_data_pt([1, 2])
    km.add_data_pt([10, 12])
    km.add_data_pt([0, -1])
    km.add_data_pt([11, 13])
    km.add_data_pt([0, 0])
    km.add_data_pt([1, 0])
    km.add_data_pt([9.4, 10])
    km.add_data_pt([9.2, 11])
    km.add_data_pt([0, 0])

IMP.set_log_level(IMP.PROGRESS)
km = IMP.kmeans.KMeans()
add_points(km)
# cluster data
k = 2
km.execute(k)
# print centers one by one
print("\nCenters:")
nCenters = km.get_number_of_centers()  # this should actually be equal to k
print("Number of centers (k): ", nCenters)
for i in range(0, nCenters):
    print("Center ", i, "   ", km.get_center(i))

# print data point assignments
n = km.get_number_of_data_points()
assignments = km.get_assignments()
sqrdist = km.get_squared_distance_to_centers()
print("Number of points: ", n)
for i in range(0, n):
    print("Point ", i, ": ")
    print(km.get_data_point(i))
    print("Cluster center: ", assignments[i], \
        " sqrt-dist: ", sqrdist[i])