58 lines
1.8 KiB
Haskell
58 lines
1.8 KiB
Haskell
{-# LANGUAGE Strict,FlexibleInstances #-}
|
|
module KMeans where
|
|
|
|
import Control.Applicative
|
|
import Control.Monad.Random
|
|
import Data.List (minimumBy, genericLength, transpose)
|
|
import Data.Ord (comparing)
|
|
import qualified Data.Map.Strict as M
|
|
|
|
|
|
type Vec = [Float]
|
|
type Cluster = [Vec]
|
|
|
|
kMeansIteration :: [Vec] -> [Vec] -> [Cluster]
|
|
kMeansIteration pts = clusterize . fixPoint iteration
|
|
where
|
|
iteration = map centroid . clusterize
|
|
|
|
clusterize centroids = M.elems $ foldr add m0 pts
|
|
where add x = M.insertWith (++) (centroids `nearestTo` x) [x]
|
|
m0 = M.unions $ map (`M.singleton` []) centroids
|
|
|
|
nearestTo :: [Vec] -> Vec -> Vec
|
|
nearestTo pts x = minimumBy (comparing (distance x)) pts
|
|
|
|
distance :: Vec -> Vec -> Float
|
|
distance a b = sum $ map (^2) $ zipWith (-) a b
|
|
|
|
centroid :: [Vec] -> Vec
|
|
centroid = map mean . transpose
|
|
where mean pts = sum pts / genericLength pts
|
|
|
|
fixPoint :: Eq a => (a -> a) -> a -> a
|
|
fixPoint f x = if x == fx then x else fixPoint f fx where fx = f x
|
|
|
|
-- initial sampling
|
|
|
|
kMeans :: MonadRandom m => Int -> [Vec] -> m [Cluster]
|
|
kMeans n pts = kMeansIteration pts <$> take n <$> randomElements pts
|
|
|
|
kMeansPP :: MonadRandom m => Int -> [Vec] -> m [Cluster]
|
|
kMeansPP n pts = kMeansIteration pts <$> centroids
|
|
where centroids = iterate (>>= nextCentroid) x0 !! (n-1)
|
|
x0 = take 1 <$> randomElements pts
|
|
nextCentroid cs = (: cs) <$> fromList (map (weight cs) pts)
|
|
weight cs x = (x, toRational $ distance x (cs `nearestTo` x))
|
|
|
|
randomElements :: MonadRandom m => [a] -> m [a]
|
|
randomElements pts = map (pts !!) <$> getRandomRs (0, length pts)
|
|
|
|
-- sample cluster generation
|
|
|
|
instance (RandomGen g, Monoid m) => Monoid (Rand g m) where
|
|
mempty = pure mempty
|
|
mappend = liftA2 mappend
|
|
|
|
mkCluster n s m = take n . transpose <$> mapM randomsAround m
|
|
where randomsAround x0 = map (\x -> x0+s*atanh x) <$> getRandomRs (-1,1)
|