{- |
Module      : ConClusion.Chemistry.Topology
Description : Principal Component Analysis
Copyright   : Phillip Seeber, 2023
License     : AGPL-3
Maintainer  : phillip.seeber@googlemail.com
Stability   : experimental
Portability : POSIX, Windows

This module implements routines to work with simple molden style XYZ trajectories. This includes
parsers as well as functions to obtain structural features in internal coordinates.

For an introduction into PCA see <https://www.cs.cmu.edu/~elaw/papers/pca.pdf>.

Diherdrals require a special metric, see <https://onlinelibrary.wiley.com/doi/full/10.1002/prot.20310)>.
-}
module ConClusion.Chemistry.Topology (
  Molecule,
  Trajectory,
  xyz,
  trajectory,
  B (..),
  A (..),
  D (..),
  Feature (..),
  getFeatures,
)
where

import ConClusion.Array.Conversion
import ConClusion.Array.Util hiding (angle)
import qualified ConClusion.Array.Util as ArrayUtil
import Data.Attoparsec.Text hiding (D)
import Data.Foldable
import Data.Massiv.Array as Massiv hiding (B, D, forM)
import qualified Numeric.LinearAlgebra as LA
import RIO
import qualified RIO.Seq as Seq
import qualified RIO.Text as Text
import qualified RIO.Vector.Boxed as VectorB

-- | A Molecule in cartesian coordinates.
data Molecule = Molecule
  { Molecule -> Double
energy :: !Double
  -- ^ The energy of the molecule.
  , Molecule -> Vector Text
atoms :: !(VectorB.Vector Text)
  -- ^ Chemical symbols or names of the atoms. \(N\) vector.
  , Molecule -> Matrix S Double
coordinates :: !(Massiv.Matrix S Double)
  -- ^ Cartesian coordinates. Atoms as rows, xyz as columns. \(N \times 3\) matrix.
  }

type Trajectory = Seq Molecule

-- | Parser for molecules in Molden XYZ format.
xyz :: Parser Molecule
xyz :: Parser Molecule
xyz = do
  -- Starting line with number of atoms.
  Int
nAtoms <- do
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Int
n <- forall a. Integral a => Parser a
decimal
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Parser Text ()
endOfLine
    forall (m :: * -> *) a. Monad m => a -> m a
return Int
n

  -- Comment line. Must contain the energy.
  Double
energy <- do
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Double
e <- Parser Text Double
double
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Parser Text ()
endOfLine
    forall (m :: * -> *) a. Monad m => a -> m a
return Double
e

  -- Multiple lines of atoms.
  [(Text, Array D Ix2 Double)]
atoms <- forall (m :: * -> *) a. Monad m => Int -> m a -> m [a]
count Int
nAtoms forall a b. (a -> b) -> a -> b
$ do
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    [Char]
name <- forall (f :: * -> *) a. Alternative f => f a -> f [a]
many1 forall a b. (a -> b) -> a -> b
$ Parser Char
letter forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser Char
digit
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Double
x <- Parser Text Double
double
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Double
y <- Parser Text Double
double
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Double
z <- Parser Text Double
double
    ()
_ <- forall (f :: * -> *) a. Alternative f => f a -> f ()
skipMany forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Parser Char
char forall a b. (a -> b) -> a -> b
$ Char
' '
    Parser Text ()
endOfLine
    let row :: Array D Ix2 Double
row = forall r ix a b.
(Index ix, Index (Lower ix), Manifest r a) =>
Sz1 -> (a -> Int -> b) -> Array r (Lower ix) a -> Array D ix b
Massiv.expandOuter (forall ix. Index ix => ix -> Sz ix
Sz Int
1) forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r e. Manifest r e => Comp -> [e] -> Vector r e
Massiv.fromList @S Comp
Seq forall a b. (a -> b) -> a -> b
$ [Double
x, Double
y, Double
z]
    forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> Text
Text.pack [Char]
name, Array D Ix2 Double
row)

  forall (m :: * -> *) a. Monad m => a -> m a
return
    Molecule
      { $sel:energy:Molecule :: Double
energy = Double
energy
      , $sel:atoms:Molecule :: Vector Text
atoms = forall a. Int -> [a] -> Vector a
VectorB.fromListN Int
nAtoms forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ [(Text, Array D Ix2 Double)]
atoms
      , $sel:coordinates:Molecule :: Matrix S Double
coordinates = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
Massiv.compute forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) r ix e.
(HasCallStack, Foldable f, Index ix, Source r e) =>
Dim -> f (Array r ix e) -> Array DL ix e
Massiv.concat' (Int -> Dim
Dim Int
2) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall a b. (a -> b) -> a -> b
$ [(Text, Array D Ix2 Double)]
atoms
      }

-- | Parser for trajectories in XYZ format as produced by CREST.
{-# SCC trajectory #-}
trajectory :: Parser Trajectory
trajectory :: Parser Trajectory
trajectory = do
  Trajectory
mols <- forall a. [a] -> Seq a
Seq.fromList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a. Alternative f => f a -> f [a]
many1 Parser Molecule
xyz
  Parser Text ()
skipSpace
  forall t. Chunk t => Parser t ()
endOfInput
  forall (m :: * -> *) a. Monad m => a -> m a
return Trajectory
mols

-- | Selection of a bond between two atoms.
data B = B Int Int

-- | Selection of an angle between three atoms.
data A = A Int Int Int

-- | Selection of a dihedral angle between four atoms. Rotation around the central two.
data D = D Int Int Int Int

-- | Selections
data Feature
  = Energy
  | Bond B
  | Angle A
  | Dihedral D

-- | Calculate a bond distance.
bond :: (MonadThrow m) => B -> Molecule -> m Double
bond :: forall (m :: * -> *). MonadThrow m => B -> Molecule -> m Double
bond (B Int
a Int
b) Molecule{Matrix S Double
coordinates :: Matrix S Double
$sel:coordinates:Molecule :: Molecule -> Matrix S Double
coordinates}
  | Int
a forall a. Eq a => a -> a -> Bool
== Int
b = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> IndexException
IndexException forall a b. (a -> b) -> a -> b
$ [Char]
"selected atoms are identical"
  | Bool
otherwise = do
      Array U Int Double
coordA <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
a)
      Array U Int Double
coordB <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
b)
      Array U Int Double
vecAB <- Array U Int Double
coordA forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.-. Array U Int Double
coordB
      forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Floating a => a -> a
sqrt forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e. (Index ix, Source r e, Num e) => Array r ix e -> e
Massiv.sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
Massiv.map (forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int)) forall a b. (a -> b) -> a -> b
$ Array U Int Double
vecAB

-- | Calculates the sinus of an angle defined by three atoms.
angle :: (MonadThrow m) => A -> Molecule -> m Double
angle :: forall (m :: * -> *). MonadThrow m => A -> Molecule -> m Double
angle (A Int
a Int
b Int
c) Molecule{Matrix S Double
coordinates :: Matrix S Double
$sel:coordinates:Molecule :: Molecule -> Matrix S Double
coordinates}
  | Int
a forall a. Eq a => a -> a -> Bool
== Int
b Bool -> Bool -> Bool
|| Int
b forall a. Eq a => a -> a -> Bool
== Int
c Bool -> Bool -> Bool
|| Int
a forall a. Eq a => a -> a -> Bool
== Int
c = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> IndexException
IndexException forall a b. (a -> b) -> a -> b
$ [Char]
"selected atoms are identical"
  | Bool
otherwise = do
      Array U Int Double
coordA <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
a)
      Array U Int Double
coordB <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
b)
      Array U Int Double
coordC <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
c)
      Array U Int Double
vecAB <- Array U Int Double
coordA forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.-. Array U Int Double
coordB
      Array U Int Double
vecCB <- Array U Int Double
coordC forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.-. Array U Int Double
coordB
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall r e.
(Numeric r e, Source r e, Floating e) =>
Vector r e -> Vector r e -> e
ArrayUtil.angle Array U Int Double
vecAB Array U Int Double
vecCB

{- | Calculates the dihedral angle defined by four atoms. Respects rotation direction. Obtains the
result in radian.
-}
dihedral' :: (MonadThrow m) => D -> Molecule -> m Double
dihedral' :: forall (m :: * -> *). MonadThrow m => D -> Molecule -> m Double
dihedral' (D Int
a Int
b Int
c Int
d) Molecule{Matrix S Double
coordinates :: Matrix S Double
$sel:coordinates:Molecule :: Molecule -> Matrix S Double
coordinates}
  | Int
a forall a. Eq a => a -> a -> Bool
== Int
b Bool -> Bool -> Bool
|| Int
a forall a. Eq a => a -> a -> Bool
== Int
c Bool -> Bool -> Bool
|| Int
a forall a. Eq a => a -> a -> Bool
== Int
d Bool -> Bool -> Bool
|| Int
b forall a. Eq a => a -> a -> Bool
== Int
c Bool -> Bool -> Bool
|| Int
b forall a. Eq a => a -> a -> Bool
== Int
d Bool -> Bool -> Bool
|| Int
c forall a. Eq a => a -> a -> Bool
== Int
d = forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> IndexException
IndexException forall a b. (a -> b) -> a -> b
$ [Char]
"selected atoms are indentical"
  | Bool
otherwise = do
      Array U Int Double
coordA <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
a)
      Array U Int Double
coordB <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
b)
      Array U Int Double
coordC <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
c)
      Array U Int Double
coordD <- forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Matrix S Double
coordinates forall r ix e (m :: * -> *).
(MonadThrow m, Index ix, Index (Lower ix), Source r e) =>
Array r ix e -> Int -> m (Array r (Lower ix) e)
!?> Int
d)
      Array U Int Double
vecAB <- Array U Int Double
coordA forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.-. Array U Int Double
coordB
      Array U Int Double
vecBC <- Array U Int Double
coordB forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.-. Array U Int Double
coordC
      Array U Int Double
vecCD <- Array U Int Double
coordC forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.-. Array U Int Double
coordD
      let planeABC :: Array U Int Double
planeABC = forall e r.
(Element e, Manifest r e, Load r Int e) =>
Vector e -> Vector r e
vecH2M forall a b. (a -> b) -> a -> b
$ forall t. Product t => Vector t -> Vector t -> Vector t
LA.cross (forall r e.
(Manifest r e, Load r Int e, Element e) =>
Vector r e -> Vector e
vecM2H Array U Int Double
vecAB) (forall r e.
(Manifest r e, Load r Int e, Element e) =>
Vector r e -> Vector e
vecM2H Array U Int Double
vecBC) :: Massiv.Vector U Double
          planeBCD :: Array U Int Double
planeBCD = forall e r.
(Element e, Manifest r e, Load r Int e) =>
Vector e -> Vector r e
vecH2M forall a b. (a -> b) -> a -> b
$ forall t. Product t => Vector t -> Vector t -> Vector t
LA.cross (forall r e.
(Manifest r e, Load r Int e, Element e) =>
Vector r e -> Vector e
vecM2H Array U Int Double
vecBC) (forall r e.
(Manifest r e, Load r Int e, Element e) =>
Vector r e -> Vector e
vecM2H Array U Int Double
vecCD) :: Massiv.Vector U Double
          normVecRot :: Vector Double
normVecRot = forall t. Product t => Vector t -> Vector t -> Vector t
LA.cross (forall r e.
(Manifest r e, Load r Int e, Element e) =>
Vector r e -> Vector e
vecM2H Array U Int Double
vecCD) (forall r e.
(Manifest r e, Load r Int e, Element e) =>
Vector r e -> Vector e
vecM2H Array U Int Double
vecBC) :: LA.Vector Double
          rotDir :: Double
rotDir =
            if forall e r.
(Element e, Manifest r e, Load r Int e) =>
Vector e -> Vector r e
vecH2M Vector Double
normVecRot forall r e.
(Numeric r e, Source r e) =>
Vector r e -> Vector r e -> e
!.! Array U Int Double
vecAB forall a. Ord a => a -> a -> Bool
< Double
0
              then -Double
1
              else Double
1
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Double
rotDir forall a. Num a => a -> a -> a
* forall r e.
(Numeric r e, Source r e, Floating e) =>
Vector r e -> Vector r e -> e
ArrayUtil.angle Array U Int Double
planeABC Array U Int Double
planeBCD

{- | Calculates a metric value of the dihedral angle defined by four atoms. This must create 2
values in the feature matrix, instead of one.
See <https://onlinelibrary.wiley.com/doi/full/10.1002/prot.20310)>
-}
dihedral :: (MonadThrow m) => D -> Molecule -> m (Double, Double)
dihedral :: forall (m :: * -> *).
MonadThrow m =>
D -> Molecule -> m (Double, Double)
dihedral D
d Molecule
mol = do
  Double
dihedRad <- forall (m :: * -> *). MonadThrow m => D -> Molecule -> m Double
dihedral' D
d Molecule
mol
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. Floating a => a -> a
sin Double
dihedRad, forall a. Floating a => a -> a
cos Double
dihedRad)

-- | Get all selected features from a molecule.
getFeature ::
  ( Traversable t
  , MonadThrow m
  ) =>
  -- | Selection of multiple features.
  t Feature ->
  -- | A Molecule where to apply to.
  Molecule ->
  m (Massiv.Vector Massiv.DL Double)
getFeature :: forall (t :: * -> *) (m :: * -> *).
(Traversable t, MonadThrow m) =>
t Feature -> Molecule -> m (Vector DL Double)
getFeature t Feature
sel Molecule
mol = do
  t (Either Double (Double, Double))
features <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for t Feature
sel forall a b. (a -> b) -> a -> b
$ \case
    Feature
Energy -> forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> Either a b
Left forall b c a. (b -> c) -> (a -> b) -> a -> c
. Molecule -> Double
energy forall a b. (a -> b) -> a -> b
$ Molecule
mol
    Bond B
b -> forall a b. a -> Either a b
Left forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadThrow m => B -> Molecule -> m Double
bond B
b Molecule
mol
    Angle A
a -> forall a b. a -> Either a b
Left forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadThrow m => A -> Molecule -> m Double
angle A
a Molecule
mol
    Dihedral D
d -> forall a b. b -> Either a b
Right forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadThrow m =>
D -> Molecule -> m (Double, Double)
dihedral D
d Molecule
mol

  let featureVec :: Vector DL Double
featureVec =
        forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
          ( \Vector DL Double
acc Either Double (Double, Double)
val -> case Either Double (Double, Double)
val of
              Left Double
s -> Vector DL Double
acc forall r e.
(Size r, Load r Int e) =>
Vector r e -> e -> Vector DL e
`Massiv.snoc` Double
s
              Right (Double
a, Double
b) -> Vector DL Double
acc forall r e.
(Size r, Load r Int e) =>
Vector r e -> e -> Vector DL e
`Massiv.snoc` Double
a forall r e.
(Size r, Load r Int e) =>
Vector r e -> e -> Vector DL e
`Massiv.snoc` Double
b
          )
          forall a. Monoid a => a
mempty
          t (Either Double (Double, Double))
features

  forall (m :: * -> *) a. Monad m => a -> m a
return Vector DL Double
featureVec

{- | Obtains the feature matrix \(\mathbf{X}\) for a principal component analysis. Given \(m\)
features to analyse in \(n\) measurements, \(\mathbf{X}\) will be a \(m \times n\) matrix.
-}
{-# SCC getFeatures #-}
getFeatures ::
  ( MonadThrow m
  , Traversable f
  ) =>
  f Feature ->
  Trajectory ->
  m (Massiv.Matrix DL Double)
getFeatures :: forall (m :: * -> *) (f :: * -> *).
(MonadThrow m, Traversable f) =>
f Feature -> Trajectory -> m (Array DL Ix2 Double)
getFeatures f Feature
sel Trajectory
trj = forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Molecule -> m (Array D Ix2 Double)
toCols Trajectory
trj forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall r ix e (f :: * -> *) (m :: * -> *).
(MonadThrow m, Foldable f, Index ix, Source r e) =>
Dim -> f (Array r ix e) -> m (Array DL ix e)
concatM (Int -> Dim
Dim Int
1)
 where
  toCols :: Molecule -> m (Array D Ix2 Double)
toCols Molecule
v = forall r ix a b.
(Index ix, Index (Lower ix), Manifest r a) =>
Sz1 -> (a -> Int -> b) -> Array r (Lower ix) a -> Array D ix b
expandInner @U @Ix2 (forall ix. Index ix => ix -> Sz ix
Sz Int
1) forall a b. a -> b -> a
const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute @U forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *).
(Traversable t, MonadThrow m) =>
t Feature -> Molecule -> m (Vector DL Double)
getFeature f Feature
sel Molecule
v