DBSCAN.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- | Density-based clustering (DBSCAN) with brute-force region queries.
module DataFrame.DBSCAN (
    module DataFrame.Model,
    DBSCANConfig (..),
    defaultDBSCANConfig,
    DBSCANModel (..),
    dbscanSurrogateExpr,
) where

import Control.Monad.ST (runST)
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Unboxed.Mutable as VUM

import DataFrame.DecisionTree.Fit (fitDecisionTree)
import DataFrame.DecisionTree.Types (TreeConfig)
import DataFrame.Featurize.Internal (
    Features (..),
    columnExprName,
    extractFeatures,
    materializeColumn,
 )
import qualified DataFrame.Functions as F
import qualified DataFrame.Internal.Column as DI
import DataFrame.Internal.DataFrame (DataFrame, fromColumns)
import DataFrame.Internal.Expression (Expr)
import DataFrame.LinearAlgebra (epsNeighbors)
import DataFrame.Model

data DBSCANConfig = DBSCANConfig
    { dbEps :: !Double
    , dbMinSamples :: !Int
    }
    deriving (Eq, Show)

defaultDBSCANConfig :: DBSCANConfig
defaultDBSCANConfig = DBSCANConfig{dbEps = 0.5, dbMinSamples = 5}

-- | A fitted DBSCAN labelling, with -1 marking noise.
data DBSCANModel = DBSCANModel
    { dbLabels :: !(VU.Vector Int)
    , dbCoreSampleIndices :: !(VU.Vector Int)
    , dbNClusters :: !Int
    }
    deriving (Eq, Show)

instance Fit DBSCANConfig [Expr Double] where
    type ModelOf DBSCANConfig [Expr Double] = DBSCANModel
    fit = fitDBSCAN

-- | Cluster the feature columns with DBSCAN.
fitDBSCAN :: DBSCANConfig -> [Expr Double] -> DataFrame -> DBSCANModel
fitDBSCAN cfg features df =
    DBSCANModel labels coreIdx nClusters
  where
    Features _ _ rows n _ = extractFeatures features df
    nbrs = V.generate n (epsNeighbors (dbEps cfg) rows)
    isCore i = VU.length (nbrs V.! i) + 1 >= dbMinSamples cfg
    coreIdx = VU.fromList [i | i <- [0 .. n - 1], isCore i]
    labels = clusterLabels n nbrs isCore
    nClusters = if VU.null labels then 0 else 1 + maximum (-1 : VU.toList labels)

clusterLabels ::
    Int -> V.Vector (VU.Vector Int) -> (Int -> Bool) -> VU.Vector Int
clusterLabels n nbrs isCore = runST $ do
    lab <- VUM.replicate n (-2)
    let seedLoop c i
            | i >= n = pure ()
            | otherwise = do
                li <- VUM.read lab i
                if li /= -2
                    then seedLoop c (i + 1)
                    else
                        if not (isCore i)
                            then VUM.write lab i (-1) >> seedLoop c (i + 1)
                            else do
                                VUM.write lab i c
                                expand lab c (VU.toList (nbrs V.! i))
                                seedLoop (c + 1) (i + 1)
        expand _ _ [] = pure ()
        expand lab c (q : qs) = do
            lq <- VUM.read lab q
            if lq == -1
                then VUM.write lab q c >> expand lab c qs
                else
                    if lq /= -2
                        then expand lab c qs
                        else do
                            VUM.write lab q c
                            let extra = if isCore q then VU.toList (nbrs V.! q) else []
                            expand lab c (extra ++ qs)
    seedLoop 0 0
    VU.freeze lab

-- | Fit a decision-tree surrogate on the DBSCAN labels.
dbscanSurrogateExpr ::
    TreeConfig -> [Expr Double] -> DBSCANModel -> DataFrame -> Expr Int
dbscanSurrogateExpr cfg features model df =
    fitDecisionTree cfg (F.col @Int clusterCol) augmented
  where
    clusterCol = "__cluster__"
    cols = map (\e -> (columnExprName e, materializeColumn df e)) features
    augmented =
        fromColumns $
            [(n, DI.fromUnboxedVector v) | (n, v) <- cols]
                ++ [(clusterCol, DI.fromUnboxedVector (dbLabels model))]