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
data Molecule = Molecule
{ Molecule -> Double
energy :: !Double
, Molecule -> Vector Text
atoms :: !(VectorB.Vector Text)
, Molecule -> Matrix S Double
coordinates :: !(Massiv.Matrix S Double)
}
type Trajectory = Seq Molecule
xyz :: Parser Molecule
xyz :: Parser Molecule
xyz = do
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
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
[(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
}
{-# 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
data B = B Int Int
data A = A Int Int Int
data D = D Int Int Int Int
data Feature
= Energy
| Bond B
| Angle A
| Dihedral D
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
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
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
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)
getFeature ::
( Traversable t
, MonadThrow m
) =>
t Feature ->
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
{-# 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