module AI.Clustering.KMeans
( KMeans(..)
, kmeans
, kmeansBy
, kmeansWith
, Method(..)
, decode
, withinSS
) where
import Control.Monad (forM_)
import Control.Monad.Primitive (PrimMonad, PrimState)
import qualified Data.Matrix.Unboxed as MU
import qualified Data.Matrix.Generic as MG
import qualified Data.Matrix.Unboxed.Mutable as MM
import Data.Ord (comparing)
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Mutable as VM
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM
import Data.List (minimumBy, foldl')
import System.Random.MWC (Gen)
import AI.Clustering.KMeans.Types (KMeans(..), Method(..))
import AI.Clustering.KMeans.Internal (sumSquares, forgy, kmeansPP)
kmeans :: (PrimMonad m, MG.Matrix mat U.Vector Double)
=> Gen (PrimState m)
-> Method
-> Int
-> mat U.Vector Double
-> m KMeans
kmeans g method k mat = kmeansBy g method k dat (MG.takeRow mat)
where
dat = U.enumFromN 0 $ MG.rows mat
kmeansBy :: (PrimMonad m, G.Vector v a)
=> Gen (PrimState m)
-> Method
-> Int
-> v a
-> (a -> U.Vector Double)
-> m KMeans
kmeansBy g method k dat fn = do
initial <- case method of
Forgy -> forgy g k dat fn
KMeansPP -> kmeansPP g k dat fn
return $ kmeansWith initial dat fn
kmeansWith :: G.Vector v a
=> MU.Matrix Double
-> v a
-> (a -> U.Vector Double)
-> KMeans
kmeansWith initial dat fn | d /= MU.cols initial || k > n = error "check input"
| otherwise = KMeans member centers
where
(member, centers) = loop initial U.empty
loop means membership
| membership' == membership = (membership, means)
| otherwise = loop (update membership') membership'
where
membership' = assign means
assign means = U.generate n $ \i ->
let x = fn $ G.unsafeIndex dat i
in fst $ minimumBy (comparing snd) $ zip [0..k1] $ map (sumSquares x) $ MU.toRows means
update membership = MU.create $ do
m <- MM.replicate (k,d) 0.0
count <- UM.replicate k (0 :: Int)
forM_ [0..n1] $ \i -> do
let x = membership U.! i
UM.unsafeRead count x >>= UM.unsafeWrite count x . (+1)
let vec = fn $ dat G.! i
forM_ [0..d1] $ \j ->
MM.unsafeRead m (x,j) >>= MM.unsafeWrite m (x,j) . (+ (vec U.! j))
forM_ [0..k1] $ \i -> do
c <- UM.unsafeRead count i
forM_ [0..d1] $ \j ->
MM.unsafeRead m (i,j) >>= MM.unsafeWrite m (i,j) . (/fromIntegral c)
return m
n = G.length dat
k = MU.rows initial
d = MU.cols initial
decode :: KMeans -> [a] -> [[a]]
decode result xs = V.toList $ V.create $ do
v <- VM.replicate n []
forM_ (zip (U.toList membership) xs) $ \(i,x) ->
VM.unsafeRead v i >>= VM.unsafeWrite v i . (x:)
return v
where
membership = _clusters result
n = U.maximum membership + 1
withinSS :: KMeans -> MU.Matrix Double -> [Double]
withinSS result mat = zipWith f (decode result [0 .. MU.rows mat1]) .
MU.toRows . _centers $ result
where
f c center = foldl' (+) 0 $ map (sumSquares center . MU.takeRow mat) c