-- |
-- Module      :  Statistics.Covariance.LedoitWolf
-- Description :  Shrinkage based covariance estimator by Ledoit and Wolf
-- Copyright   :  (c) 2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Creation date: Thu Sep  9 14:08:26 2021.
module Statistics.Covariance.LedoitWolf
  ( ledoitWolf,
  )
where

import Data.Foldable
import qualified Numeric.LinearAlgebra as L
import Statistics.Covariance.Internal.Tools
import Statistics.Covariance.Types

-- | Shrinkage based covariance estimator by Ledoit and Wolf.
--
-- See Ledoit, O., & Wolf, M., A well-conditioned estimator for
-- large-dimensional covariance matrices, Journal of Multivariate Analysis,
-- 88(2), 365–411 (2004). http://dx.doi.org/10.1016/s0047-259x(03)00096-4.
--
-- Return 'Left' if
--
-- - only one sample is available.
--
-- - no parameters are available.
--
-- NOTE: This function may call 'error' due to partial library functions.
ledoitWolf ::
  DoCenter ->
  -- | Sample data matrix of dimension \(n \times p\), where \(n\) is the number
  -- of samples (rows), and \(p\) is the number of parameters (columns).
  L.Matrix Double ->
  Either String (L.Herm Double)
ledoitWolf :: DoCenter -> Matrix Double -> Either String (Herm Double)
ledoitWolf DoCenter
c Matrix Double
xs
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = String -> Either String (Herm Double)
forall a b. a -> Either a b
Left String
"ledoitWolf: Need more than one sample."
  | Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = String -> Either String (Herm Double)
forall a b. a -> Either a b
Left String
"ledoitWolf: Need at least one parameter."
  -- The Ledoit and Wolf shrinkage estimator of the covariance matrix
  -- (Equation 14). However, a different, more general formula avoiding a2 is
  -- used. See Equation (4) in Chen2010b.
  | Bool
otherwise = Herm Double -> Either String (Herm Double)
forall a b. b -> Either a b
Right (Herm Double -> Either String (Herm Double))
-> Herm Double -> Either String (Herm Double)
forall a b. (a -> b) -> a -> b
$ Double -> Herm Double -> Double -> Herm Double -> Herm Double
shrinkWith Double
rho Herm Double
sigma Double
mu Herm Double
im
  where
    n :: Int
n = Matrix Double -> Int
forall t. Matrix t -> Int
L.rows Matrix Double
xs
    p :: Int
p = Matrix Double -> Int
forall t. Matrix t -> Int
L.cols Matrix Double
xs
    (Vector Double
means, Herm Double
sigma) = Matrix Double -> (Vector Double, Herm Double)
L.meanCov Matrix Double
xs
    xsCentered :: Matrix Double
xsCentered = case DoCenter
c of
      DoCenter
DoCenter -> Vector Double -> Matrix Double -> Matrix Double
centerWith Vector Double
means Matrix Double
xs
      DoCenter
AssumeCentered -> Matrix Double
xs
    im :: Herm Double
im = Matrix Double -> Herm Double
forall t. Matrix t -> Herm t
L.trustSym (Matrix Double -> Herm Double) -> Matrix Double -> Herm Double
forall a b. (a -> b) -> a -> b
$ Int -> Matrix Double
forall a. (Num a, Element a) => Int -> Matrix a
L.ident Int
p
    mu :: Double
mu = Herm Double -> Double
muE Herm Double
sigma
    d2 :: Double
d2 = Herm Double -> Herm Double -> Double -> Double
d2E Herm Double
im Herm Double
sigma Double
mu
    b2 :: Double
b2 = Matrix Double -> Herm Double -> Double -> Double
b2E Matrix Double
xsCentered Herm Double
sigma Double
d2
    rho :: Double
rho = Double
b2 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d2

-- Inner product for symmetric matrices based on an adjusted Frobenius norm (p
-- 376).
--
-- NOTE: This function is commutative (and therefor qualified as an inner
-- product) for symmetric matrices only, and not in the general case.
frobenius :: L.Matrix Double -> L.Matrix Double -> Double
frobenius :: Matrix Double -> Matrix Double -> Double
frobenius Matrix Double
xs Matrix Double
ys
  | Int
xsRows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
xsCols = String -> Double
forall a. HasCallStack => String -> a
error String
"frobenius: Bug! Left matrix is not square."
  | Int
ysRows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
ysCols = String -> Double
forall a. HasCallStack => String -> a
error String
"frobenius: Bug! Right matrix is not square."
  | Int
xsRows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
ysRows = String -> Double
forall a. HasCallStack => String -> a
error String
"frobenius: Bug! Matrices have different size."
  | Bool
otherwise = Double -> Double
forall a. Fractional a => a -> a
recip (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
xsRows) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Matrix Double -> Double
trace (Matrix Double
xs Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
L.<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
L.tr' Matrix Double
ys)
  where
    xsRows :: Int
xsRows = Matrix Double -> Int
forall t. Matrix t -> Int
L.rows Matrix Double
xs
    xsCols :: Int
xsCols = Matrix Double -> Int
forall t. Matrix t -> Int
L.cols Matrix Double
xs
    ysRows :: Int
ysRows = Matrix Double -> Int
forall t. Matrix t -> Int
L.rows Matrix Double
ys
    ysCols :: Int
ysCols = Matrix Double -> Int
forall t. Matrix t -> Int
L.cols Matrix Double
ys

-- Estimator of mu (Lemma 3.2).
muE ::
  -- Sample covariance matrix.
  L.Herm Double ->
  Double
-- Avoid matrix multiplication in Frobenius norm.
muE :: Herm Double -> Double
muE Herm Double
sigma' = Double -> Double
forall a. Fractional a => a -> a
recip Double
xsRows Double -> Double -> Double
forall a. Num a => a -> a -> a
* Matrix Double -> Double
trace Matrix Double
sigma
  where
    sigma :: Matrix Double
sigma = Herm Double -> Matrix Double
forall t. Herm t -> Matrix t
L.unSym Herm Double
sigma'
    xsRows :: Double
xsRows = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ Matrix Double -> Int
forall t. Matrix t -> Int
L.rows Matrix Double
sigma

-- Estimator of d2 (Lemma 3.3).
d2E ::
  -- Identity matrix.
  L.Herm Double ->
  -- Sample covariance matrix.
  L.Herm Double ->
  -- Estimate of mu.
  Double ->
  Double
d2E :: Herm Double -> Herm Double -> Double -> Double
d2E Herm Double
im Herm Double
sigma Double
mu = Matrix Double -> Matrix Double -> Double
frobenius Matrix Double
m Matrix Double
m
  where
    m :: Matrix Double
m = Herm Double -> Matrix Double
forall t. Herm t -> Matrix t
L.unSym Herm Double
sigma Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
- Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
L.scale Double
mu (Herm Double -> Matrix Double
forall t. Herm t -> Matrix t
L.unSym Herm Double
im)

-- Estimator of b2 (Lemma 3.4).
b2E ::
  -- Data matrix.
  L.Matrix Double ->
  -- Sample covariance matrix.
  L.Herm Double ->
  -- Estimate of d2.
  Double ->
  Double
b2E :: Matrix Double -> Herm Double -> Double -> Double
b2E Matrix Double
xs Herm Double
sigma = Double -> Double -> Double
forall a. Ord a => a -> a -> a
min Double
b2m
  where
    n :: Double
n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ Matrix Double -> Int
forall t. Matrix t -> Int
L.rows Matrix Double
xs
    -- NOTE: The authors use a transposed data matrix. They refer to one out of
    -- n columns, each with p rows. Here, we have one out of n rows, each with p
    -- columns.
    --
    -- Long story short. Each y is an observation, a vector of length p.
    ys :: [Double]
ys =
      [ Matrix Double -> Matrix Double -> Double
frobenius Matrix Double
d Matrix Double
d
        | Vector Double
y <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
L.toRows Matrix Double
xs,
          let d :: Matrix Double
d = (Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
L.asColumn Vector Double
y Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
L.<> Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
L.asRow Vector Double
y) Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
- Herm Double -> Matrix Double
forall t. Herm t -> Matrix t
L.unSym Herm Double
sigma
      ]
    b2m :: Double
b2m = Double -> Double
forall a. Fractional a => a -> a
recip (Double
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
n) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double -> Double -> Double) -> Double -> [Double] -> Double
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0 [Double]
ys