-- | Linear array ops for PCA: center, Gram (dagger-compose), BLAS mult.
--
-- Data layout is always /samples × features/ (rows = observations).
--
-- The dagger promise used here is the array transpose: covariance on the
-- feature side is @X† X@ via transpose-compose (up to scaling). Spectral
-- work is delegated to hmatrix (same bridge pattern as @Harpie.Hmatrix@,
-- which is flag-gated in harpie).
module Circuit.PCA.Lin
  ( -- * Layout
    nSamples,
    nFeatures,

    -- * Center
    columnMeans,
    centerColumns,

    -- * Dagger-compose
    gramFeatures,

    -- * BLAS
    multM,
    transpose2,

    -- * hmatrix bridge
    toMatrix,
    fromMatrix,
  )
where

import Data.Vector.Storable qualified as VS
import Data.Vector.Unboxed qualified as VU
import Harpie.Array.Storable (Array)
import Harpie.Array.Storable qualified as A
import Numeric.LinearAlgebra (Matrix, cols, flatten, reshape, rows, (<>))
import Numeric.LinearAlgebra qualified as LA
import Prelude hiding ((<>))

-- | Number of samples (rows).
nSamples :: Array Double -> Int
nSamples :: Array Double -> Int
nSamples Array Double
a = case Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array Double -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
A.shape Array Double
a) of
  (Int
n : [Int]
_) -> Int
n
  [Int]
_ -> Int
0

-- | Number of features (columns).
nFeatures :: Array Double -> Int
nFeatures :: Array Double -> Int
nFeatures Array Double
a = case Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array Double -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
A.shape Array Double
a) of
  [Int
_, Int
p] -> Int
p
  [Int
p] -> Int
p
  [Int]
_ -> Int
0

-- | Rank-2 storable array → hmatrix (Nothing if not rank 2).
toMatrix :: Array Double -> Maybe (Matrix Double)
toMatrix :: Array Double -> Maybe (Matrix Double)
toMatrix Array Double
a =
  case Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
VU.toList (Array Double -> Vector Int
forall {k} (v :: k -> *) (a :: k). Array v a -> Vector Int
A.shape Array Double
a) of
    [Int
_, Int
c] -> Matrix Double -> Maybe (Matrix Double)
forall a. a -> Maybe a
Just (Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape Int
c (Array Double -> Vector Double
forall t a. FromVector t a => t -> Vector a
A.asVector Array Double
a))
    [Int]
_ -> Maybe (Matrix Double)
forall a. Maybe a
Nothing

-- | hmatrix → rank-2 storable array.
fromMatrix :: Matrix Double -> Array Double
fromMatrix :: Matrix Double -> Array Double
fromMatrix Matrix Double
m = [Int] -> Vector Double -> Array Double
forall t a. FromVector t a => [Int] -> t -> Array a
A.array [Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
m, Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
m] (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
flatten Matrix Double
m)

-- | BLAS matrix multiply for rank-2 arrays.
multM :: Array Double -> Array Double -> Array Double
multM :: Array Double -> Array Double -> Array Double
multM Array Double
a Array Double
b =
  case (Array Double -> Maybe (Matrix Double)
toMatrix Array Double
a, Array Double -> Maybe (Matrix Double)
toMatrix Array Double
b) of
    (Just Matrix Double
ma, Just Matrix Double
mb) -> Matrix Double -> Array Double
fromMatrix (Matrix Double
ma Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
mb)
    (Maybe (Matrix Double), Maybe (Matrix Double))
_ -> [Char] -> Array Double
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.PCA.Lin.multM: expected rank-2 arrays"

-- | Column means as a length-@p@ vector (shape @[p]@).
columnMeans :: Array Double -> Array Double
columnMeans :: Array Double -> Array Double
columnMeans Array Double
x =
  case Array Double -> Maybe (Matrix Double)
toMatrix Array Double
x of
    Maybe (Matrix Double)
Nothing -> [Char] -> Array Double
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.PCA.Lin.columnMeans: expected rank-2 samples×features"
    Just Matrix Double
m ->
      let n :: Double
n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
m) :: Double
       in Matrix Double -> Double -> Array Double
columnMeansM Matrix Double
m Double
n

columnMeansM :: Matrix Double -> Double -> Array Double
columnMeansM :: Matrix Double -> Double -> Array Double
columnMeansM Matrix Double
m Double
n =
  let p :: Int
p = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
m
      means :: Vector Double
means =
        Int -> (Int -> Double) -> Vector Double
forall a. Storable a => Int -> (Int -> a) -> Vector a
VS.generate Int
p ((Int -> Double) -> Vector Double)
-> (Int -> Double) -> Vector Double
forall a b. (a -> b) -> a -> b
$ \Int
j ->
          Matrix Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
LA.sumElements (Matrix Double
m Matrix Double -> [Int] -> Matrix Double
forall t. Element t => Matrix t -> [Int] -> Matrix t
LA.¿ [Int
j]) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n
   in [Int] -> Vector Double -> Array Double
forall t a. FromVector t a => [Int] -> t -> Array a
A.array [Int
p] Vector Double
means

-- | Center columns; returns @(centered, means)@.
centerColumns :: Array Double -> (Array Double, Array Double)
centerColumns :: Array Double -> (Array Double, Array Double)
centerColumns Array Double
x =
  case Array Double -> Maybe (Matrix Double)
toMatrix Array Double
x of
    Maybe (Matrix Double)
Nothing -> [Char] -> (Array Double, Array Double)
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.PCA.Lin.centerColumns: expected rank-2 samples×features"
    Just Matrix Double
m ->
      let n :: Double
n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
m) :: Double
          mu :: Array Double
mu = Matrix Double -> Double -> Array Double
columnMeansM Matrix Double
m Double
n
          muFlat :: Vector Double
muFlat = Array Double -> Vector Double
forall t a. FromVector t a => t -> Vector a
A.asVector Array Double
mu
          xc :: Matrix Double
xc = Matrix Double
m Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
- Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asRow Vector Double
muFlat
       in (Matrix Double -> Array Double
fromMatrix Matrix Double
xc, Array Double
mu)

-- | Feature Gram matrix @X† X@ (features × features), unnormalised.
--
-- Dagger-compose: reverse wire is transpose, compose is BLAS mult.
-- Scale by @1/(n-1)@ for sample covariance if desired.
gramFeatures :: Array Double -> Array Double
gramFeatures :: Array Double -> Array Double
gramFeatures Array Double
x =
  case Array Double -> Maybe (Matrix Double)
toMatrix Array Double
x of
    Maybe (Matrix Double)
Nothing -> [Char] -> Array Double
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.PCA.Lin.gramFeatures: expected rank-2"
    Just Matrix Double
m -> Matrix Double -> Array Double
fromMatrix (Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
LA.tr Matrix Double
m Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix Double
m)

-- | Rank-2 transpose.
transpose2 :: Array Double -> Array Double
transpose2 :: Array Double -> Array Double
transpose2 Array Double
a =
  case Array Double -> Maybe (Matrix Double)
toMatrix Array Double
a of
    Maybe (Matrix Double)
Nothing -> [Char] -> Array Double
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.PCA.Lin.transpose2: expected rank-2"
    Just Matrix Double
m -> Matrix Double -> Array Double
fromMatrix (Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
LA.tr Matrix Double
m)