Topology.hs

{- |
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
  { energy :: !Double
  -- ^ The energy of the molecule.
  , atoms :: !(VectorB.Vector Text)
  -- ^ Chemical symbols or names of the atoms. \(N\) vector.
  , 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 = do
  -- Starting line with number of atoms.
  nAtoms <- do
    _ <- skipMany . char $ ' '
    n <- decimal
    _ <- skipMany . char $ ' '
    endOfLine
    return n

  -- Comment line. Must contain the energy.
  energy <- do
    _ <- skipMany . char $ ' '
    e <- double
    _ <- skipMany . char $ ' '
    endOfLine
    return e

  -- Multiple lines of atoms.
  atoms <- count nAtoms $ do
    _ <- skipMany . char $ ' '
    name <- many1 $ letter <|> digit
    _ <- skipMany1 . char $ ' '
    x <- double
    _ <- skipMany1 . char $ ' '
    y <- double
    _ <- skipMany1 . char $ ' '
    z <- double
    _ <- skipMany . char $ ' '
    endOfLine
    let row = Massiv.expandOuter (Sz 1) const . Massiv.fromList @S Seq $ [x, y, z]
    return (Text.pack name, row)

  return
    Molecule
      { energy = energy
      , atoms = VectorB.fromListN nAtoms . fmap fst $ atoms
      , coordinates = Massiv.compute . Massiv.concat' (Dim 2) . fmap snd $ atoms
      }

-- | Parser for trajectories in XYZ format as produced by CREST.
{-# SCC trajectory #-}
trajectory :: Parser Trajectory
trajectory = do
  mols <- Seq.fromList <$> many1 xyz
  skipSpace
  endOfInput
  return 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 (B a b) Molecule{coordinates}
  | a == b = throwM . IndexException $ "selected atoms are identical"
  | otherwise = do
      coordA <- compute @U <$> (coordinates !?> a)
      coordB <- compute @U <$> (coordinates !?> b)
      vecAB <- coordA .-. coordB
      return . sqrt . Massiv.sum . Massiv.map (^ (2 :: Int)) $ vecAB

-- | Calculates the sinus of an angle defined by three atoms.
angle :: (MonadThrow m) => A -> Molecule -> m Double
angle (A a b c) Molecule{coordinates}
  | a == b || b == c || a == c = throwM . IndexException $ "selected atoms are identical"
  | otherwise = do
      coordA <- compute @U <$> (coordinates !?> a)
      coordB <- compute @U <$> (coordinates !?> b)
      coordC <- compute @U <$> (coordinates !?> c)
      vecAB <- coordA .-. coordB
      vecCB <- coordC .-. coordB
      return $ ArrayUtil.angle vecAB 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' (D a b c d) Molecule{coordinates}
  | a == b || a == c || a == d || b == c || b == d || c == d = throwM . IndexException $ "selected atoms are indentical"
  | otherwise = do
      coordA <- compute @U <$> (coordinates !?> a)
      coordB <- compute @U <$> (coordinates !?> b)
      coordC <- compute @U <$> (coordinates !?> c)
      coordD <- compute @U <$> (coordinates !?> d)
      vecAB <- coordA .-. coordB
      vecBC <- coordB .-. coordC
      vecCD <- coordC .-. coordD
      let planeABC = vecH2M $ LA.cross (vecM2H vecAB) (vecM2H vecBC) :: Massiv.Vector U Double
          planeBCD = vecH2M $ LA.cross (vecM2H vecBC) (vecM2H vecCD) :: Massiv.Vector U Double
          normVecRot = LA.cross (vecM2H vecCD) (vecM2H vecBC) :: LA.Vector Double
          rotDir =
            if vecH2M normVecRot !.! vecAB < 0
              then -1
              else 1
      return $ rotDir * ArrayUtil.angle planeABC 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 d mol = do
  dihedRad <- dihedral' d mol
  return (sin dihedRad, cos 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 sel mol = do
  features <- for sel $ \case
    Energy -> return . Left . energy $ mol
    Bond b -> Left <$> bond b mol
    Angle a -> Left <$> angle a mol
    Dihedral d -> Right <$> dihedral d mol

  let featureVec =
        foldl
          ( \acc val -> case val of
              Left s -> acc `Massiv.snoc` s
              Right (a, b) -> acc `Massiv.snoc` a `Massiv.snoc` b
          )
          mempty
          features

  return 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 sel trj = traverse toCols trj >>= concatM (Dim 1)
 where
  toCols v = expandInner @U @Ix2 (Sz 1) const . compute @U <$> getFeature sel v