Matrix.hs

-- |
-- Module      :  ELynx.Tools.Matrix
-- Copyright   :  (c) Dominik Schrempf 2020
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Feb 14 13:33:13 2019.
--
-- Tools for matrices from 'Numeric.LinearAlgebra'.
module ELynx.Tools.Matrix
  ( -- * Matrices
    fMapRowSeq,
    fMapRowPar,
    fMapRowParChunk,
    fMapColSeq,
    fMapColPar,
    fMapColParChunk,
    (|||),
    (===),
    subSampleMatrix,
  )
where

import Control.Parallel.Strategies
import Data.List
import qualified Data.Matrix.Generic as M
import qualified Data.Vector.Generic as V

-- | Map a function on each row of a DIM2 array; sequential version.
fMapRowSeq ::
  (V.Vector v a, V.Vector v b) =>
  (v a -> v b) ->
  M.Matrix v a ->
  M.Matrix v b
fMapRowSeq f m = M.fromRows $ map f (M.toRows m)

-- | Map a function on each row of a DIM2 array; parallel version.
fMapRowPar ::
  (V.Vector v a, V.Vector v b) =>
  (v a -> v b) ->
  M.Matrix v a ->
  M.Matrix v b
fMapRowPar f m = M.fromRows $ parMap rseq f (M.toRows m)

-- | Map a function on each row of a DIM2 array; parallel version with given chunk size.
fMapRowParChunk ::
  (V.Vector v a, V.Vector v b) =>
  Int ->
  (v a -> v b) ->
  M.Matrix v a ->
  M.Matrix v b
fMapRowParChunk n f m =
  M.fromRows (map f (M.toRows m) `using` parListChunk n rseq)

-- | Map a function on each row of a DIM2 array; sequential version.
fMapColSeq ::
  (V.Vector v a, V.Vector v b) =>
  (v a -> v b) ->
  M.Matrix v a ->
  M.Matrix v b
fMapColSeq f m = M.fromColumns $ map f (M.toColumns m)

-- | Map a function on each row of a DIM2 array; parallel version.
fMapColPar ::
  (V.Vector v a, V.Vector v b) =>
  (v a -> v b) ->
  M.Matrix v a ->
  M.Matrix v b
fMapColPar f m = M.fromColumns $ parMap rseq f (M.toColumns m)

-- | Map a function on each row of a DIM2 array; parallel version with given chunk size.
fMapColParChunk ::
  (V.Vector v a, V.Vector v b) =>
  Int ->
  (v a -> v b) ->
  M.Matrix v a ->
  M.Matrix v b
fMapColParChunk n f m =
  M.fromColumns (map f (M.toColumns m) `using` parListChunk n rseq)

-- | Horizontal concatenation.
(|||) :: (V.Vector v a) => M.Matrix v a -> M.Matrix v a -> M.Matrix v a
(|||) l r = M.fromColumns $ lCs ++ rCs
  where
    lCs = M.toColumns l
    rCs = M.toColumns r

-- | Vertical concatenation.
(===) :: (V.Vector v a) => M.Matrix v a -> M.Matrix v a -> M.Matrix v a
(===) l r = M.fromRows $ lRs ++ rRs
  where
    lRs = M.toRows l
    rRs = M.toRows r

-- | Sample the given sites from a matrix.
subSampleMatrix :: (V.Vector v a) => [Int] -> M.Matrix v a -> M.Matrix v a
subSampleMatrix is m =
  M.fromColumns $ foldl' (\a i -> M.takeColumn m i : a) [] is