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

{- | Principal component analysis via the symmetric Jacobi eigensolver on the
covariance of the (optionally standardized) feature columns. 'fit' trains a
'PCAModel' (components + explained variance); the projection is exposed as
'pcaExprs' / 'pcaTransform' (PCA is a transformer, so it has no 'Predict').
-}
module DataFrame.PCA (
    module DataFrame.Model,
    NComponents (..),
    PCAConfig (..),
    defaultPCAConfig,
    PCAModel (..),
    pcaExprs,
    pcaTransform,
) 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 (gram)
import DataFrame.LinearAlgebra.Eigen (jacobiEigenSym)
import DataFrame.Model
import DataFrame.Operators ((.*.), (.+.), (.-.))
import DataFrame.Transform (Transform (..))

-- | How many components to keep.
data NComponents = NComp !Int | VarianceCovered !Double
    deriving (NComponents -> NComponents -> Bool
(NComponents -> NComponents -> Bool)
-> (NComponents -> NComponents -> Bool) -> Eq NComponents
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NComponents -> NComponents -> Bool
== :: NComponents -> NComponents -> Bool
$c/= :: NComponents -> NComponents -> Bool
/= :: NComponents -> NComponents -> Bool
Eq, Int -> NComponents -> ShowS
[NComponents] -> ShowS
NComponents -> String
(Int -> NComponents -> ShowS)
-> (NComponents -> String)
-> ([NComponents] -> ShowS)
-> Show NComponents
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NComponents -> ShowS
showsPrec :: Int -> NComponents -> ShowS
$cshow :: NComponents -> String
show :: NComponents -> String
$cshowList :: [NComponents] -> ShowS
showList :: [NComponents] -> ShowS
Show)

data PCAConfig = PCAConfig
    { PCAConfig -> NComponents
pcaNComponents :: !NComponents
    , PCAConfig -> Bool
pcaStandardize :: !Bool
    }
    deriving (PCAConfig -> PCAConfig -> Bool
(PCAConfig -> PCAConfig -> Bool)
-> (PCAConfig -> PCAConfig -> Bool) -> Eq PCAConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PCAConfig -> PCAConfig -> Bool
== :: PCAConfig -> PCAConfig -> Bool
$c/= :: PCAConfig -> PCAConfig -> Bool
/= :: PCAConfig -> PCAConfig -> Bool
Eq, Int -> PCAConfig -> ShowS
[PCAConfig] -> ShowS
PCAConfig -> String
(Int -> PCAConfig -> ShowS)
-> (PCAConfig -> String)
-> ([PCAConfig] -> ShowS)
-> Show PCAConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PCAConfig -> ShowS
showsPrec :: Int -> PCAConfig -> ShowS
$cshow :: PCAConfig -> String
show :: PCAConfig -> String
$cshowList :: [PCAConfig] -> ShowS
showList :: [PCAConfig] -> ShowS
Show)

defaultPCAConfig :: PCAConfig
defaultPCAConfig :: PCAConfig
defaultPCAConfig = PCAConfig{pcaNComponents :: NComponents
pcaNComponents = Int -> NComponents
NComp Int
2, pcaStandardize :: Bool
pcaStandardize = Bool
False}

{- | A fitted PCA. 'pcaComponents' are sklearn's @components_@ (row @i@ is the
@i@-th loading vector); 'pcaScale' is @Just@ the per-column std when
standardizing.
-}
data PCAModel = PCAModel
    { PCAModel -> Vector (Vector Double)
pcaComponents :: !(V.Vector (VU.Vector Double))
    , PCAModel -> Vector Double
pcaExplainedVariance :: !(VU.Vector Double)
    , PCAModel -> Vector Double
pcaExplainedVarianceRatio :: !(VU.Vector Double)
    , PCAModel -> Vector Double
pcaMean :: !(VU.Vector Double)
    , PCAModel -> Maybe (Vector Double)
pcaScale :: !(Maybe (VU.Vector Double))
    , PCAModel -> Vector Text
pcaFeatureNames :: !(V.Vector T.Text)
    }
    deriving (PCAModel -> PCAModel -> Bool
(PCAModel -> PCAModel -> Bool)
-> (PCAModel -> PCAModel -> Bool) -> Eq PCAModel
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PCAModel -> PCAModel -> Bool
== :: PCAModel -> PCAModel -> Bool
$c/= :: PCAModel -> PCAModel -> Bool
/= :: PCAModel -> PCAModel -> Bool
Eq, Int -> PCAModel -> ShowS
[PCAModel] -> ShowS
PCAModel -> String
(Int -> PCAModel -> ShowS)
-> (PCAModel -> String) -> ([PCAModel] -> ShowS) -> Show PCAModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PCAModel -> ShowS
showsPrec :: Int -> PCAModel -> ShowS
$cshow :: PCAModel -> String
show :: PCAModel -> String
$cshowList :: [PCAModel] -> ShowS
showList :: [PCAModel] -> ShowS
Show)

instance Fit PCAConfig [Expr Double] where
    type ModelOf PCAConfig [Expr Double] = PCAModel
    fit :: CheckFrame
  (FrameReq PCAConfig [Expr Double]) (FrameFor [Expr Double]) =>
PCAConfig
-> [Expr Double]
-> FrameFor [Expr Double]
-> FitResult
     (FrameFor [Expr Double]) (ModelOf PCAConfig [Expr Double])
fit = PCAConfig -> [Expr Double] -> DataFrame -> PCAModel
PCAConfig
-> [Expr Double]
-> FrameFor [Expr Double]
-> FitResult
     (FrameFor [Expr Double]) (ModelOf PCAConfig [Expr Double])
fitPCA

-- | Fit PCA on the given feature columns (each must be a @Col@).
fitPCA :: PCAConfig -> [Expr Double] -> DataFrame -> PCAModel
fitPCA :: PCAConfig -> [Expr Double] -> DataFrame -> PCAModel
fitPCA PCAConfig
cfg [Expr Double]
features DataFrame
df =
    PCAModel
        { pcaComponents :: Vector (Vector Double)
pcaComponents = Int -> Vector (Vector Double) -> Vector (Vector Double)
forall a. Int -> Vector a -> Vector a
V.take Int
k Vector (Vector Double)
vecs
        , pcaExplainedVariance :: Vector Double
pcaExplainedVariance = Int -> Vector Double -> Vector Double
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
k Vector Double
evar
        , pcaExplainedVarianceRatio :: Vector Double
pcaExplainedVarianceRatio = Int -> Vector Double -> Vector Double
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
k Vector Double
ratio
        , pcaMean :: Vector Double
pcaMean = Vector Double
means
        , pcaScale :: Maybe (Vector Double)
pcaScale = if PCAConfig -> Bool
pcaStandardize PCAConfig
cfg then Vector Double -> Maybe (Vector Double)
forall a. a -> Maybe a
Just Vector Double
scales else Maybe (Vector Double)
forall a. Maybe a
Nothing
        , pcaFeatureNames :: Vector Text
pcaFeatureNames = [Text] -> Vector Text
forall a. [a] -> Vector a
V.fromList [Text]
names
        }
  where
    Features [Text]
names [Vector Double]
cols Vector (Vector Double)
_ Int
n Int
d = [Expr Double] -> DataFrame -> Features
extractFeatures [Expr Double]
features DataFrame
df
    means :: Vector Double
means = [Double] -> Vector Double
forall a. Unbox a => [a] -> Vector a
VU.fromList [Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector Double
c 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
n) | Vector Double
c <- [Vector Double]
cols]
    scales :: Vector Double
scales =
        [Double] -> Vector Double
forall a. Unbox a => [a] -> Vector a
VU.fromList
            [ let mu :: Double
mu = Vector Double
means Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j
                  v :: Double
v = Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum ((Double -> Double) -> Vector Double -> Vector Double
forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
VU.map (\Double
x -> (Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) Vector Double
c) 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
n)
                  s :: Double
s = Double -> Double
forall a. Floating a => a -> a
sqrt Double
v
               in if Double
s Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
1e-12 then Double
1 else Double
s
            | (Int
j, Vector Double
c) <- [Int] -> [Vector Double] -> [(Int, Vector Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Vector Double]
cols
            ]
    scaled :: Vector (Vector Double)
scaled =
        Int -> (Int -> Vector Double) -> Vector (Vector Double)
forall a. Int -> (Int -> a) -> Vector a
V.generate Int
n ((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
d ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
j ->
                let mu :: Double
mu = Vector Double
means Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j
                    s :: Double
s = if PCAConfig -> Bool
pcaStandardize PCAConfig
cfg then Vector Double
scales Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
j else Double
1
                 in (([Vector Double]
cols [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
j) Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
mu) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s
    denom :: Double
denom = 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
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
    cov :: Vector (Vector Double)
cov = (Vector Double -> Vector Double)
-> Vector (Vector Double) -> Vector (Vector Double)
forall a b. (a -> b) -> Vector a -> Vector b
V.map ((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
denom)) (Vector (Vector Double) -> Vector (Vector Double)
gram Vector (Vector Double)
scaled)
    (Vector Double
evals, Vector (Vector Double)
vecs) = Vector (Vector Double) -> (Vector Double, Vector (Vector Double))
jacobiEigenSym Vector (Vector Double)
cov
    evar :: Vector Double
evar = (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. Ord a => a -> a -> a
max Double
0) Vector Double
evals
    total :: Double
total = Vector Double -> Double
forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector Double
evar
    ratio :: Vector Double
ratio = if Double
total Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Vector Double
evar else (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
total) Vector Double
evar
    k :: Int
k = NComponents -> Int -> Vector Double -> Int
resolveK (PCAConfig -> NComponents
pcaNComponents PCAConfig
cfg) Int
d Vector Double
ratio

-- | Per-component projection expressions, named @pc1@, @pc2@, …
pcaExprs :: PCAModel -> [(T.Text, Expr Double)]
pcaExprs :: PCAModel -> [(Text, Expr Double)]
pcaExprs PCAModel
m =
    [ (Text
"pc" 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), Vector Double -> Expr Double
componentExpr (PCAModel -> Vector (Vector Double)
pcaComponents PCAModel
m Vector (Vector Double) -> Int -> Vector Double
forall a. Vector a -> Int -> a
V.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)))
    | Int
i <- [Int
1 .. Vector (Vector Double) -> Int
forall a. Vector a -> Int
V.length (PCAModel -> Vector (Vector Double)
pcaComponents PCAModel
m)]
    ]
  where
    names :: [Text]
names = Vector Text -> [Text]
forall a. Vector a -> [a]
V.toList (PCAModel -> Vector Text
pcaFeatureNames PCAModel
m)
    means :: [Double]
means = Vector Double -> [Double]
forall a. Unbox a => Vector a -> [a]
VU.toList (PCAModel -> Vector Double
pcaMean PCAModel
m)
    scales :: [Double]
scales = [Double]
-> (Vector Double -> [Double]) -> Maybe (Vector Double) -> [Double]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Double -> [Double]
forall a. a -> [a]
repeat Double
1) Vector Double -> [Double]
forall a. Unbox a => Vector a -> [a]
VU.toList (PCAModel -> Maybe (Vector Double)
pcaScale PCAModel
m)
    componentExpr :: Vector Double -> Expr Double
componentExpr Vector Double
vec =
        (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
$
            [ Double -> Expr Double
forall a. Columnable a => a -> Expr a
F.lit (Double
w Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
s) Expr Double -> Expr Double -> Expr Double
forall a. (Columnable a, Num a) => Expr a -> Expr a -> Expr a
.*. ((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
mu)
            | (Double
w, Text
n, Double
mu, Double
s) <- [Double]
-> [Text]
-> [Double]
-> [Double]
-> [(Double, Text, Double, Double)]
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 (Vector Double -> [Double]
forall a. Unbox a => Vector a -> [a]
VU.toList Vector Double
vec) [Text]
names [Double]
means [Double]
scales
            ]

-- | The PCA projection as a composable fitted 'Transform'.
pcaTransform :: PCAModel -> Transform
pcaTransform :: PCAModel -> Transform
pcaTransform PCAModel
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) <- PCAModel -> [(Text, Expr Double)]
pcaExprs PCAModel
m]

resolveK :: NComponents -> Int -> VU.Vector Double -> Int
resolveK :: NComponents -> Int -> Vector Double -> Int
resolveK (NComp Int
k) Int
d Vector Double
_ = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k Int
d)
resolveK (VarianceCovered Double
frac) Int
d Vector Double
ratio = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
d (Int -> Double -> Int -> Int
go Int
0 Double
0 Int
1))
  where
    go :: Int -> Double -> Int -> Int
go !Int
acc !Double
cum !Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
ratio = Vector Double -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Double
ratio
        | Double
cum Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
frac = Int
acc
        | Bool
otherwise = Int -> Double -> Int -> Int
go (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Double
cum Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Vector Double
ratio Vector Double -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
VU.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

zip4 :: [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 :: forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 (a
a : [a]
as) (b
b : [b]
bs) (c
c : [c]
cs) (d
d : [d]
ds) = (a
a, b
b, c
c, d
d) (a, b, c, d) -> [(a, b, c, d)] -> [(a, b, c, d)]
forall a. a -> [a] -> [a]
: [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
forall a b c d. [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]
zip4 [a]
as [b]
bs [c]
cs [d]
ds
zip4 [a]
_ [b]
_ [c]
_ [d]
_ = []