RFF.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

{- | Approximate RBF-kernel SVM via Random Fourier Features (Rahimi & Recht): map
each row through @z(x) = √(2/D)·cos(W x + b)@ with @W ~ N(0, 2γI)@ (seeded), then
fit a linear SVC in the random-feature space. 'predict' compiles to a closed
@Σ_r β_r·cos(…)@ expression of size @O(D·d)@, independent of the row count.
-}
module DataFrame.SVM.RFF (
    RFFConfig (..),
    defaultRFFConfig,
    RFFSVMModel (..),
) where

import Data.List (sort)
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU

import DataFrame.Featurize.Internal (featureNames, numericMatrix, targetValues)
import qualified DataFrame.Functions as F
import DataFrame.Internal.Column (Columnable)
import DataFrame.Internal.DataFrame (DataFrame)
import DataFrame.Internal.Expression (Expr (..))
import DataFrame.LinearAlgebra (dot)
import DataFrame.LinearSolver (LinearModel (..), SolverConfig (..), fitProx)
import DataFrame.LinearSolver.Loss (sqHingeLoss)
import DataFrame.Model (Fit (..), Predict (..))
import DataFrame.Operators ((.*.), (.+.), (.>.))
import DataFrame.Random (Gen, gaussianVector, mkGen, nextDouble)

data RFFConfig = RFFConfig
    { rffD :: !Int
    , rffGamma :: !Double
    , rffC :: !Double
    , rffMaxIter :: !Int
    , rffTol :: !Double
    , rffSeed :: !Int
    }
    deriving (Eq, Show)

defaultRFFConfig :: RFFConfig
defaultRFFConfig =
    RFFConfig
        { rffD = 100
        , rffGamma = 0.1
        , rffC = 1.0
        , rffMaxIter = 1000
        , rffTol = 1.0e-4
        , rffSeed = 0
        }

{- | A fitted RFF SVM (binary). 'rffW' / 'rffB' are the random projection;
'rffCoef' / 'rffIntercept' the linear SVC in feature space.
-}
data RFFSVMModel a = RFFSVMModel
    { rffW :: !(V.Vector (VU.Vector Double))
    , rffB :: !(VU.Vector Double)
    , rffCoef :: !(VU.Vector Double)
    , rffIntercept :: !Double
    , rffScale :: !Double
    , rffNegClass :: !a
    , rffPosClass :: !a
    , rffFeatureNames :: !(V.Vector T.Text)
    }
    deriving (Show)

instance (Columnable a, Ord a) => Fit RFFConfig (Expr a) (RFFSVMModel a) where
    fit = fitRFFSVM

instance (Columnable a) => Predict (RFFSVMModel a) a where
    predict m =
        If (margin .>. F.lit 0) (Lit (rffPosClass m)) (Lit (rffNegClass m))
      where
        names = V.toList (rffFeatureNames m)
        margin =
            foldr (.+.) (F.lit (rffIntercept m)) $
                [ F.lit (rffCoef m VU.! r * rffScale m) .*. cosTerm r
                | r <- [0 .. V.length (rffW m) - 1]
                , rffCoef m VU.! r /= 0
                ]
        cosTerm r = cos (linComb (rffW m V.! r) (rffB m VU.! r))
        linComb w b =
            foldr (.+.) (F.lit b) $
                [ F.lit (w VU.! j) .*. (Col n :: Expr Double)
                | (j, n) <- zip [0 ..] names
                ]

-- | Fit a binary RFF SVM. Targets with more than two classes are rejected.
fitRFFSVM ::
    (Columnable a, Ord a) => RFFConfig -> Expr a -> DataFrame -> RFFSVMModel a
fitRFFSVM cfg target df =
    case classes of
        [neg, pos] -> build neg pos
        _ -> error "fitRFFSVM: binary classification only (got /= 2 classes)"
  where
    names = featureNames target df
    (nameVec, mat) = numericMatrix names df
    ys = targetValues target df
    classes = sort (foldr dedup [] (V.toList ys))
    dedup x acc = if x `elem` acc then acc else x : acc
    d = if V.null mat then 0 else VU.length (V.head mat)
    bigD = max 1 (rffD cfg)
    (ws, bs) = sampleRFF bigD d (rffGamma cfg) (mkGen (rffSeed cfg))
    scale = sqrt (2 / fromIntegral bigD)
    z = V.map (featureRow ws bs scale) mat
    featNames = V.fromList ["rff" <> T.pack (show r) | r <- [0 .. bigD - 1]]
    build neg pos =
        let labels = VU.generate (V.length ys) (\i -> if ys V.! i == pos then 1 else -1)
            solverCfg =
                SolverConfig
                    { scL1Lambda = 0
                    , scL2Lambda = 1 / rffC cfg
                    , scMaxIter = rffMaxIter cfg
                    , scTol = rffTol cfg
                    , scSampleWeights = Nothing
                    }
            model = fitProx sqHingeLoss solverCfg z labels featNames
         in RFFSVMModel ws bs (lmWeights model) (lmIntercept model) scale neg pos nameVec

sampleRFF ::
    Int -> Int -> Double -> Gen -> (V.Vector (VU.Vector Double), VU.Vector Double)
sampleRFF bigD d gamma g0 = (V.fromList ws, VU.fromList bs)
  where
    sigma = sqrt (2 * gamma)
    (ws, g1) = goW bigD g0 []
    goW 0 g acc = (reverse acc, g)
    goW k g acc =
        let (vec, g') = gaussianVector d g
         in goW (k - 1) g' (VU.map (* sigma) vec : acc)
    bs = take bigD (goB g1)
    goB g = let (u, g') = nextDouble g in (u * 2 * pi) : goB g'

featureRow ::
    V.Vector (VU.Vector Double) ->
    VU.Vector Double ->
    Double ->
    VU.Vector Double ->
    VU.Vector Double
featureRow ws bs scale x =
    VU.generate (V.length ws) $ \r ->
        scale * cos (dot (ws V.! r) x + bs VU.! r)