<s> import numpy as np <EOL> import time <EOL> import parakeet <EOL> from parakeet import allpairs , each <EOL> from parakeet . testing_helpers import eq , run_local_tests <EOL> def python_update_assignments ( X , centroids ) : <EOL> dists = [ [ np . sum ( ( x - y ) ** <NUM_LIT:2> ) for y in centroids ] for x in X ] <EOL> return np . argmin ( dists , <NUM_LIT:1> ) <EOL> def python_update_centroids ( X , assignments , k ) : <EOL> d = X . shape [ <NUM_LIT:1> ] <EOL> new_centroids = np . zeros ( ( k , d ) , dtype = X . dtype ) <EOL> for i in xrange ( k ) : <EOL> mask = assignments == i <EOL> count = np . sum ( mask ) <EOL> assigned_data = X [ mask ] <EOL> if count > <NUM_LIT:1> : <EOL> new_centroids [ i , : ] = np . mean ( assigned_data ) <EOL> elif count == <NUM_LIT:1> : <EOL> new_centroids [ i , : ] = assigned_data [ <NUM_LIT:0> ] <EOL> return new_centroids <EOL> def python_kmeans ( X , k , maxiters = <NUM_LIT:100> , initial_assignments = None ) : <EOL> n = X . shape [ <NUM_LIT:0> ] <EOL> if initial_assignments is None : <EOL> assignments = np . random . randint ( <NUM_LIT:0> , k , size = n ) <EOL> else : <EOL> assignments = initial_assignments <EOL> centroids = python_update_centroids ( X , assignments , k ) <EOL> for _ in xrange ( maxiters ) : <EOL> old_assignments = assignments <EOL> assignments = python_update_assignments ( X , centroids ) <EOL> if all ( old_assignments == assignments ) : <EOL> break <EOL> centroids = python_update_centroids ( X , assignments , k ) <EOL> return centroids <EOL> def sqr_dist ( x , y ) : <EOL> return sum ( ( x - y ) ** <NUM_LIT:2> ) <EOL> def parakeet_update_assignments ( X , centroids ) : <EOL> dists = allpairs ( sqr_dist , X , centroids ) <EOL> return np . argmin ( dists , <NUM_LIT:1> ) <EOL> def mean ( X ) : <EOL> return sum ( X ) / len ( X ) <EOL> def parakeet_update_centroids ( X , assignments , k ) : <EOL> d = X . shape [ <NUM_LIT:1> ] <EOL> new_centroids = np . zeros ( ( k , d ) , dtype = X . dtype ) <EOL> for i in xrange ( k ) : <EOL> mask = ( assignments == i ) <EOL> count = np . sum ( mask ) <EOL> assigned_data = X [ mask ] <EOL> if count == <NUM_LIT:1> : <EOL> new_centroids [ i , : ] = assigned_data [ <NUM_LIT:0> ] <EOL> elif count > <NUM_LIT:1> : <EOL> new_centroids [ i , : ] = parakeet . mean ( assigned_data ) <EOL> return new_centroids <EOL> def parakeet_kmeans ( X , k , maxiters = <NUM_LIT:100> , initial_assignments = None ) : <EOL> n = X . shape [ <NUM_LIT:0> ] <EOL> if initial_assignments is None : <EOL> assignments = np . random . randint ( <NUM_LIT:0> , k , size = n ) <EOL> else : <EOL> assignments = initial_assignments <EOL> centroids = python_update_centroids ( X , assignments , k ) <EOL> for _ in xrange ( maxiters ) : <EOL> old_assignments = assignments <EOL> assignments = parakeet_update_assignments ( X , centroids ) <EOL> if all ( old_assignments == assignments ) : <EOL> break <EOL> centroids = python_update_centroids ( X , assignments , k ) <EOL> return centroids <EOL> def test_kmeans ( ) : <EOL> n = <NUM_LIT:200> <EOL> d = <NUM_LIT:4> <EOL> X = np . random . randn ( n * d ) . reshape ( n , d ) <EOL> k = <NUM_LIT:2> <EOL> niters = <NUM_LIT:10> <EOL> assignments = np . random . randint ( <NUM_LIT:0> , k , size = n ) <EOL> parakeet_C = parakeet_kmeans ( X , k , niters , assignments ) <EOL> python_C = python_kmeans ( X , k , niters , assignments ) <EOL> assert eq ( parakeet_C . shape , python_C . shape ) , "<STR_LIT>" % ( python_C . shape , parakeet_C . shape ) <EOL> assert eq ( parakeet_C , python_C ) , "<STR_LIT>" % ( python_C , parakeet_C ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> run_local_tests ( ) </s>
