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

{- | Kernel PCA with an RBF kernel, solved on a set of landmark points (Nyström).
Exact kernel PCA when the landmark count covers every row, a principled
approximation otherwise. 'fit' trains the model; the projection is exposed as
'kernelPCAExprs' / 'kernelPcaTransform' (a transformer, so no 'Predict').
-}
module DataFrame.PCA.Kernel (
    module DataFrame.Model,
    KernelPCAConfig (..),
    defaultKernelPCAConfig,
    KernelPCAModel (..),
    kernelPCAExprs,
    kernelPcaTransform,
) where

import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU

import DataFrame.Featurize.Internal (Features (..), extractFeatures)
import qualified DataFrame.Functions as F
import DataFrame.Internal.DataFrame (DataFrame)
import DataFrame.Internal.Expression (Expr (..), UExpr (..))
import DataFrame.LinearAlgebra (sqDist)
import DataFrame.LinearAlgebra.Eigen (jacobiEigenSym)
import DataFrame.Model
import DataFrame.Operators ((.*.), (.+.), (.-.))
import DataFrame.Random (mkGen, sampleIndices)
import DataFrame.Transform (Transform (..))

data KernelPCAConfig = KernelPCAConfig
    { KernelPCAConfig -> Int
kpcaNComponents :: !Int
    , KernelPCAConfig -> Maybe Double
kpcaGamma :: !(Maybe Double)
    , KernelPCAConfig -> Int
kpcaNLandmarks :: !Int
    , KernelPCAConfig -> Int
kpcaSeed :: !Int
    }
    deriving (KernelPCAConfig -> KernelPCAConfig -> Bool
(KernelPCAConfig -> KernelPCAConfig -> Bool)
-> (KernelPCAConfig -> KernelPCAConfig -> Bool)
-> Eq KernelPCAConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: KernelPCAConfig -> KernelPCAConfig -> Bool
== :: KernelPCAConfig -> KernelPCAConfig -> Bool
$c/= :: KernelPCAConfig -> KernelPCAConfig -> Bool
/= :: KernelPCAConfig -> KernelPCAConfig -> Bool
Eq, Int -> KernelPCAConfig -> ShowS
[KernelPCAConfig] -> ShowS
KernelPCAConfig -> String
(Int -> KernelPCAConfig -> ShowS)
-> (KernelPCAConfig -> String)
-> ([KernelPCAConfig] -> ShowS)
-> Show KernelPCAConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KernelPCAConfig -> ShowS
showsPrec :: Int -> KernelPCAConfig -> ShowS
$cshow :: KernelPCAConfig -> String
show :: KernelPCAConfig -> String
$cshowList :: [KernelPCAConfig] -> ShowS
showList :: [KernelPCAConfig] -> ShowS
Show)

defaultKernelPCAConfig :: KernelPCAConfig
defaultKernelPCAConfig :: KernelPCAConfig
defaultKernelPCAConfig =
    KernelPCAConfig
        { kpcaNComponents :: Int
kpcaNComponents = Int
2
        , kpcaGamma :: Maybe Double
kpcaGamma = Maybe Double
forall a. Maybe a
Nothing
        , kpcaNLandmarks :: Int
kpcaNLandmarks = Int
128
        , kpcaSeed :: Int
kpcaSeed = Int
0
        }

{- | A fitted kernel PCA. Each component is @Σ_l βₗ·K(x, landmarkₗ) + cᵢ@ with an
RBF kernel of bandwidth 'kpcaGammaUsed'.
-}
data KernelPCAModel = KernelPCAModel
    { KernelPCAModel -> Vector (Vector Double)
kpcaLandmarks :: !(V.Vector (VU.Vector Double))
    , KernelPCAModel -> Vector (Vector Double)
kpcaBetas :: !(V.Vector (VU.Vector Double))
    , KernelPCAModel -> Vector Double
kpcaConsts :: !(VU.Vector Double)
    , KernelPCAModel -> Vector Double
kpcaEigenvalues :: !(VU.Vector Double)
    , KernelPCAModel -> Double
kpcaGammaUsed :: !Double
    , KernelPCAModel -> Vector Text
kpcaFeatureNames :: !(V.Vector T.Text)
    }
    deriving (KernelPCAModel -> KernelPCAModel -> Bool
(KernelPCAModel -> KernelPCAModel -> Bool)
-> (KernelPCAModel -> KernelPCAModel -> Bool) -> Eq KernelPCAModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: KernelPCAModel -> KernelPCAModel -> Bool
== :: KernelPCAModel -> KernelPCAModel -> Bool
$c/= :: KernelPCAModel -> KernelPCAModel -> Bool
/= :: KernelPCAModel -> KernelPCAModel -> Bool
Eq, Int -> KernelPCAModel -> ShowS
[KernelPCAModel] -> ShowS
KernelPCAModel -> String
(Int -> KernelPCAModel -> ShowS)
-> (KernelPCAModel -> String)
-> ([KernelPCAModel] -> ShowS)
-> Show KernelPCAModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KernelPCAModel -> ShowS
showsPrec :: Int -> KernelPCAModel -> ShowS
$cshow :: KernelPCAModel -> String
show :: KernelPCAModel -> String
$cshowList :: [KernelPCAModel] -> ShowS
showList :: [KernelPCAModel] -> ShowS
Show)

instance Fit KernelPCAConfig [Expr Double] where
    type ModelOf KernelPCAConfig [Expr Double] = KernelPCAModel
    fit :: CheckFrame
  (FrameReq KernelPCAConfig [Expr Double])
  (FrameFor [Expr Double]) =>
KernelPCAConfig
-> [Expr Double]
-> FrameFor [Expr Double]
-> FitResult
     (FrameFor [Expr Double]) (ModelOf KernelPCAConfig [Expr Double])
fit = KernelPCAConfig -> [Expr Double] -> DataFrame -> KernelPCAModel
KernelPCAConfig
-> [Expr Double]
-> FrameFor [Expr Double]
-> FitResult
     (FrameFor [Expr Double]) (ModelOf KernelPCAConfig [Expr Double])
fitKernelPCA

-- | Fit kernel PCA over the given feature columns.
fitKernelPCA :: KernelPCAConfig -> [Expr Double] -> DataFrame -> KernelPCAModel
fitKernelPCA :: KernelPCAConfig -> [Expr Double] -> DataFrame -> KernelPCAModel
fitKernelPCA KernelPCAConfig
cfg [Expr Double]
features DataFrame
df =
    KernelPCAModel
        { kpcaLandmarks :: Vector (Vector Double)
kpcaLandmarks = Vector (Vector Double)
landmarks
        , kpcaBetas :: Vector (Vector Double)
kpcaBetas = Vector (Vector Double)
betas
        , kpcaConsts :: Vector Double
kpcaConsts = Vector Double
consts
        , kpcaEigenvalues :: Vector Double
kpcaEigenvalues = Int -> Vector Double -> Vector Double
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
k Vector Double
evals
        , kpcaGammaUsed :: Double
kpcaGammaUsed = Double
gamma
        , kpcaFeatureNames :: Vector Text
kpcaFeatureNames = [Text] -> Vector Text
forall a. [a] -> Vector a
V.fromList [Text]
names
        }
  where
    Features [Text]
names [Vector Double]
_ Vector (Vector Double)
rows Int
n Int
d = [Expr Double] -> DataFrame -> Features
extractFeatures [Expr Double]
features DataFrame
df
    m :: Int
m = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (KernelPCAConfig -> Int
kpcaNLandmarks KernelPCAConfig
cfg)) Int
n
    (Vector Int
idx, Gen
_) = Int -> Int -> Gen -> (Vector Int, Gen)
sampleIndices Int
m Int
n (Int -> Gen
mkGen (KernelPCAConfig -> Int
kpcaSeed KernelPCAConfig
cfg))
    landmarks :: Vector (Vector Double)
landmarks = (Int -> Vector Double) -> Vector Int -> Vector (Vector Double)
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Vector (Vector Double)
rows Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.!) (Vector Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert Vector Int
idx)
    gamma :: Double
gamma = case KernelPCAConfig -> Maybe Double
kpcaGamma KernelPCAConfig
cfg of
        Just Double
g -> Double
g
        Maybe Double
Nothing -> Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
d)
    kmat :: Vector (Vector Double)
kmat =
        Int -> (Int -> Vector Double) -> Vector (Vector Double)
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
m ((Int -> Vector Double) -> Vector (Vector Double))
-> (Int -> Vector Double) -> Vector (Vector Double)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            Int -> (Int -> Double) -> Vector Double
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
m ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
j ->
                Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double
forall a. Num a => a -> a
negate Double
gamma Double -> Double -> Double
forall a. Num a => a -> a -> a
* Vector Double -> Vector Double -> Double
sqDist (Vector (Vector Double)
landmarks Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i) (Vector (Vector Double)
landmarks Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
j))
    rowMean :: Int -> Double
rowMean Int
i = Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum (Vector (Vector Double)
kmat Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m
    totalMean :: Double
totalMean = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Double
rowMean Int
i | Int
i <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m
    centered :: Vector (Vector Double)
centered =
        Int -> (Int -> Vector Double) -> Vector (Vector Double)
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
m ((Int -> Vector Double) -> Vector (Vector Double))
-> (Int -> Vector Double) -> Vector (Vector Double)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            Int -> (Int -> Double) -> Vector Double
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
m ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
j ->
                (Vector (Vector Double)
kmat Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i) Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
rowMean Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
- Int -> Double
rowMean Int
j Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
totalMean
    (Vector Double
evals, Vector (Vector Double)
vecs) = Vector (Vector Double) -> (Vector Double, Vector (Vector Double))
jacobiEigenSym Vector (Vector Double)
centered
    k :: Int
k = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (KernelPCAConfig -> Int
kpcaNComponents KernelPCAConfig
cfg) Int
m
    alphas :: Vector (Vector Double)
alphas =
        Int -> (Int -> Vector Double) -> Vector (Vector Double)
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
k ((Int -> Vector Double) -> Vector (Vector Double))
-> (Int -> Vector Double) -> Vector (Vector Double)
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            let lam :: Double
lam = Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1e-12 (Vector Double
evals Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i)
             in (Double -> Double) -> Vector Double -> Vector Double
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt Double
lam) (Vector (Vector Double)
vecs Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i)
    betas :: Vector (Vector Double)
betas =
        (Vector Double -> Vector Double)
-> Vector (Vector Double) -> Vector (Vector Double)
forall a b. (a -> b) -> Vector a -> Vector b
V.map
            (\Vector Double
a -> let s :: Double
s = Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector Double
a Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
m in (Double -> Double) -> Vector Double -> Vector Double
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (Double -> Double -> Double
forall a. Num a => a -> a -> a
subtract Double
s) Vector Double
a)
            Vector (Vector Double)
alphas
    consts :: Vector Double
consts =
        Int -> (Int -> Double) -> Vector Double
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
k ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
i ->
            let a :: Vector Double
a = Vector (Vector Double)
alphas Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i
                sA :: Double
sA = Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector Double
a
             in Double -> Double
forall a. Num a => a -> a
negate ([Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Vector Double
a Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
l Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
rowMean Int
l | Int
l <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
                    Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
totalMean Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
sA

-- | Per-component projection expressions, named @kpc1@, @kpc2@, …
kernelPCAExprs :: KernelPCAModel -> [(T.Text, Expr Double)]
kernelPCAExprs :: KernelPCAModel -> [(Text, Expr Double)]
kernelPCAExprs KernelPCAModel
m =
    [ (Text
"kpc" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)), Int -> Expr Double
componentExpr Int
i)
    | Int
i <- [Int
0 .. Vector (Vector Double) -> Int
forall a. Vector a -> Int
V.length (KernelPCAModel -> Vector (Vector Double)
kpcaBetas KernelPCAModel
m) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ]
  where
    names :: [Text]
names = Vector Text -> [Text]
forall a. Vector a -> [a]
V.toList (KernelPCAModel -> Vector Text
kpcaFeatureNames KernelPCAModel
m)
    gamma :: Double
gamma = KernelPCAModel -> Double
kpcaGammaUsed KernelPCAModel
m
    componentExpr :: Int -> Expr Double
componentExpr Int
i =
        (Expr Double -> Expr Double -> Expr Double)
-> Expr Double -> [Expr Double] -> Expr Double
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Expr a -> Expr a
(.+.) (Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (KernelPCAModel -> Vector Double
kpcaConsts KernelPCAModel
m Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i)) ([Expr Double] -> Expr Double) -> [Expr Double] -> Expr Double
forall a b. (a -> b) -> a -> b
$
            [ Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (KernelPCAModel -> Vector (Vector Double)
kpcaBetas KernelPCAModel
m Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
i Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
l) Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Expr a -> Expr a
.*. Vector Double -> Expr Double
kernelExpr (KernelPCAModel -> Vector (Vector Double)
kpcaLandmarks KernelPCAModel
m Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! Int
l)
            | Int
l <- [Int
0 .. Vector (Vector Double) -> Int
forall a. Vector a -> Int
V.length (KernelPCAModel -> Vector (Vector Double)
kpcaLandmarks KernelPCAModel
m) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ]
    kernelExpr :: Vector Double -> Expr Double
kernelExpr Vector Double
landmark =
        Expr Double -> Expr Double
forall a. Floating a => a -> a
exp (Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (Double -> Double
forall a. Num a => a -> a
negate Double
gamma) Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Expr a -> Expr a
.*. Vector Double -> Expr Double
sqDistExpr Vector Double
landmark)
    sqDistExpr :: Vector Double -> Expr Double
sqDistExpr Vector Double
landmark =
        (Expr Double -> Expr Double -> Expr Double)
-> Expr Double -> [Expr Double] -> Expr Double
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Expr a -> Expr a
(.+.) (Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
0) ([Expr Double] -> Expr Double) -> [Expr Double] -> Expr Double
forall a b. (a -> b) -> a -> b
$
            [ let diff :: Expr Double
diff = (Text -> Expr Double
forall a. Columnable a => Text -> Expr a
Col Text
n :: Expr Double) Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Expr a -> Expr a
.-. Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit Double
lj in Expr Double
diff Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Expr a -> Expr a
.*. Expr Double
diff
            | (Text
n, Double
lj) <- [Text] -> [Double] -> [(Text, Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Text]
names (Vector Double -> [Double]
forall a. Unbox a => Vector a -> [a]
VU.toList Vector Double
landmark)
            ]

-- | The kernel-PCA projection as a composable fitted 'Transform'.
kernelPcaTransform :: KernelPCAModel -> Transform
kernelPcaTransform :: KernelPCAModel -> Transform
kernelPcaTransform KernelPCAModel
m = [NamedExpr] -> Transform
Transform [(Text
n, Expr Double -> UExpr
forall a. Columnable a => Expr a -> UExpr
UExpr Expr Double
e) | (Text
n, Expr Double
e) <- KernelPCAModel -> [(Text, Expr Double)]
kernelPCAExprs KernelPCAModel
m]