-- |
-- Module      :  ELynx.Data.MarkovProcess.MixtureModel
-- Description :  Mixture models are a set of substitution models with weights
-- Copyright   :  (c) Dominik Schrempf 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Tue Jan 29 19:17:40 2019.
--
-- To be imported qualified.
module ELynx.Data.MarkovProcess.MixtureModel
  ( -- * Types
    Weight,
    Component (weight, substModel),
    MixtureModel (name, alphabet, components),

    -- * Getters
    getWeights,
    getSubstitutionModels,

    -- * Building mixture models
    fromSubstitutionModels,

    -- * Transformations
    concatenate,
    scale,
    normalize,
    appendNameComponents,
  )
where

import qualified Data.Vector as V
import ELynx.Data.Alphabet.Alphabet hiding (all)
import qualified ELynx.Data.MarkovProcess.SubstitutionModel as S
import Prelude

-- | Mixture model component weight.
type Weight = Double

-- | A mixture model component has a weight and a substitution model.
data Component = Component
  { Component -> Weight
weight :: Weight,
    Component -> SubstitutionModel
substModel :: S.SubstitutionModel
  }
  deriving (Int -> Component -> ShowS
[Component] -> ShowS
Component -> String
(Int -> Component -> ShowS)
-> (Component -> String)
-> ([Component] -> ShowS)
-> Show Component
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Component] -> ShowS
$cshowList :: [Component] -> ShowS
show :: Component -> String
$cshow :: Component -> String
showsPrec :: Int -> Component -> ShowS
$cshowsPrec :: Int -> Component -> ShowS
Show, ReadPrec [Component]
ReadPrec Component
Int -> ReadS Component
ReadS [Component]
(Int -> ReadS Component)
-> ReadS [Component]
-> ReadPrec Component
-> ReadPrec [Component]
-> Read Component
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Component]
$creadListPrec :: ReadPrec [Component]
readPrec :: ReadPrec Component
$creadPrec :: ReadPrec Component
readList :: ReadS [Component]
$creadList :: ReadS [Component]
readsPrec :: Int -> ReadS Component
$creadsPrec :: Int -> ReadS Component
Read)

-- | A mixture model with its components.
data MixtureModel = MixtureModel
  { -- | Name
    MixtureModel -> String
name :: S.Name,
    MixtureModel -> Alphabet
alphabet :: Alphabet,
    MixtureModel -> Vector Component
components :: V.Vector Component
  }
  deriving (Int -> MixtureModel -> ShowS
[MixtureModel] -> ShowS
MixtureModel -> String
(Int -> MixtureModel -> ShowS)
-> (MixtureModel -> String)
-> ([MixtureModel] -> ShowS)
-> Show MixtureModel
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MixtureModel] -> ShowS
$cshowList :: [MixtureModel] -> ShowS
show :: MixtureModel -> String
$cshow :: MixtureModel -> String
showsPrec :: Int -> MixtureModel -> ShowS
$cshowsPrec :: Int -> MixtureModel -> ShowS
Show, ReadPrec [MixtureModel]
ReadPrec MixtureModel
Int -> ReadS MixtureModel
ReadS [MixtureModel]
(Int -> ReadS MixtureModel)
-> ReadS [MixtureModel]
-> ReadPrec MixtureModel
-> ReadPrec [MixtureModel]
-> Read MixtureModel
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [MixtureModel]
$creadListPrec :: ReadPrec [MixtureModel]
readPrec :: ReadPrec MixtureModel
$creadPrec :: ReadPrec MixtureModel
readList :: ReadS [MixtureModel]
$creadList :: ReadS [MixtureModel]
readsPrec :: Int -> ReadS MixtureModel
$creadsPrec :: Int -> ReadS MixtureModel
Read)

-- | Get weights.
getWeights :: MixtureModel -> V.Vector Weight
getWeights :: MixtureModel -> Vector Weight
getWeights = (Component -> Weight) -> Vector Component -> Vector Weight
forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> Weight
weight (Vector Component -> Vector Weight)
-> (MixtureModel -> Vector Component)
-> MixtureModel
-> Vector Weight
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MixtureModel -> Vector Component
components

-- | Get substitution models.
getSubstitutionModels :: MixtureModel -> V.Vector S.SubstitutionModel
getSubstitutionModels :: MixtureModel -> Vector SubstitutionModel
getSubstitutionModels = (Component -> SubstitutionModel)
-> Vector Component -> Vector SubstitutionModel
forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> SubstitutionModel
substModel (Vector Component -> Vector SubstitutionModel)
-> (MixtureModel -> Vector Component)
-> MixtureModel
-> Vector SubstitutionModel
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MixtureModel -> Vector Component
components

-- | Create a mixture model from a list of substitution models.
fromSubstitutionModels :: S.Name -> V.Vector Weight -> V.Vector S.SubstitutionModel -> MixtureModel
fromSubstitutionModels :: String -> Vector Weight -> Vector SubstitutionModel -> MixtureModel
fromSubstitutionModels String
n Vector Weight
ws Vector SubstitutionModel
sms
  | Vector Weight -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Vector Weight
ws = String -> MixtureModel
forall a. HasCallStack => String -> a
error String
"fromSubstitutionModels: No weights given."
  | Vector Weight -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Weight
ws Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector SubstitutionModel -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector SubstitutionModel
sms = String -> MixtureModel
forall a. HasCallStack => String -> a
error String
"fromSubstitutionModels: Number of weights and substitution models does not match."
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Vector Alphabet -> Bool
forall a. Eq a => Vector a -> Bool
allEqual Vector Alphabet
alphs = String -> MixtureModel
forall a. HasCallStack => String -> a
error String
"fromSubstitutionModels: alphabets of substitution models are not equal."
  | Bool
otherwise = String -> Alphabet -> Vector Component -> MixtureModel
MixtureModel String
n (Vector Alphabet -> Alphabet
forall a. Vector a -> a
V.head Vector Alphabet
alphs) Vector Component
comps
  where
    comps :: Vector Component
comps = (Weight -> SubstitutionModel -> Component)
-> Vector Weight -> Vector SubstitutionModel -> Vector Component
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Weight -> SubstitutionModel -> Component
Component Vector Weight
ws Vector SubstitutionModel
sms
    alphs :: Vector Alphabet
alphs = (SubstitutionModel -> Alphabet)
-> Vector SubstitutionModel -> Vector Alphabet
forall a b. (a -> b) -> Vector a -> Vector b
V.map SubstitutionModel -> Alphabet
S.alphabet Vector SubstitutionModel
sms
    allEqual :: Vector a -> Bool
allEqual Vector a
xs
      | Vector a -> Bool
forall a. Vector a -> Bool
V.null Vector a
xs = Bool
True
      | Bool
otherwise = (a -> Bool) -> Vector a -> Bool
forall a. (a -> Bool) -> Vector a -> Bool
V.all (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== Vector a -> a
forall a. Vector a -> a
V.head Vector a
xs) Vector a
xs

-- | Concatenate mixture models.
concatenate :: S.Name -> V.Vector MixtureModel -> MixtureModel
concatenate :: String -> Vector MixtureModel -> MixtureModel
concatenate String
n Vector MixtureModel
mms = String -> Vector Weight -> Vector SubstitutionModel -> MixtureModel
fromSubstitutionModels String
n Vector Weight
ws Vector SubstitutionModel
sms
  where
    comps :: Vector Component
comps = (MixtureModel -> Vector Component)
-> Vector MixtureModel -> Vector Component
forall a b. (a -> Vector b) -> Vector a -> Vector b
V.concatMap MixtureModel -> Vector Component
components Vector MixtureModel
mms
    ws :: Vector Weight
ws = (Component -> Weight) -> Vector Component -> Vector Weight
forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> Weight
weight Vector Component
comps
    sms :: Vector SubstitutionModel
sms = (Component -> SubstitutionModel)
-> Vector Component -> Vector SubstitutionModel
forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> SubstitutionModel
substModel Vector Component
comps

scaleComponent :: Double -> Component -> Component
scaleComponent :: Weight -> Component -> Component
scaleComponent Weight
s Component
c = Component
c {substModel :: SubstitutionModel
substModel = SubstitutionModel
s'} where s' :: SubstitutionModel
s' = Weight -> SubstitutionModel -> SubstitutionModel
S.scale Weight
s (SubstitutionModel -> SubstitutionModel)
-> SubstitutionModel -> SubstitutionModel
forall a b. (a -> b) -> a -> b
$ Component -> SubstitutionModel
substModel Component
c

-- | Scale all substitution models of the mixture model.
scale :: Double -> MixtureModel -> MixtureModel
scale :: Weight -> MixtureModel -> MixtureModel
scale Weight
s MixtureModel
m = MixtureModel
m {components :: Vector Component
components = Vector Component
cs'}
  where
    cs :: Vector Component
cs = MixtureModel -> Vector Component
components MixtureModel
m
    cs' :: Vector Component
cs' = (Component -> Component) -> Vector Component -> Vector Component
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Weight -> Component -> Component
scaleComponent Weight
s) Vector Component
cs

-- | Globally normalize a mixture model so that on average one event happens per
-- unit time.
normalize :: MixtureModel -> MixtureModel
normalize :: MixtureModel -> MixtureModel
normalize MixtureModel
mm = Weight -> MixtureModel -> MixtureModel
scale (Weight
1 Weight -> Weight -> Weight
forall a. Fractional a => a -> a -> a
/ Weight
c) MixtureModel
mm
  where
    c :: Weight
c = Vector Weight -> Weight
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Vector Weight -> Weight) -> Vector Weight -> Weight
forall a b. (a -> b) -> a -> b
$ (Weight -> Weight -> Weight)
-> Vector Weight -> Vector Weight -> Vector Weight
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Weight -> Weight -> Weight
forall a. Num a => a -> a -> a
(*) Vector Weight
weights Vector Weight
scales
    weights :: Vector Weight
weights = MixtureModel -> Vector Weight
getWeights MixtureModel
mm
    scales :: Vector Weight
scales = (SubstitutionModel -> Weight)
-> Vector SubstitutionModel -> Vector Weight
forall a b. (a -> b) -> Vector a -> Vector b
V.map SubstitutionModel -> Weight
S.totalRate (Vector SubstitutionModel -> Vector Weight)
-> Vector SubstitutionModel -> Vector Weight
forall a b. (a -> b) -> a -> b
$ MixtureModel -> Vector SubstitutionModel
getSubstitutionModels MixtureModel
mm

appendNameComponent :: S.Name -> Component -> Component
appendNameComponent :: String -> Component -> Component
appendNameComponent String
n Component
c = Component
c {substModel :: SubstitutionModel
substModel = SubstitutionModel
s'}
  where
    s' :: SubstitutionModel
s' = String -> SubstitutionModel -> SubstitutionModel
S.appendName String
n (SubstitutionModel -> SubstitutionModel)
-> SubstitutionModel -> SubstitutionModel
forall a b. (a -> b) -> a -> b
$ Component -> SubstitutionModel
substModel Component
c

-- | Append byte string to all substitution models of mixture model.
appendNameComponents :: S.Name -> MixtureModel -> MixtureModel
appendNameComponents :: String -> MixtureModel -> MixtureModel
appendNameComponents String
n MixtureModel
m = MixtureModel
m {components :: Vector Component
components = Vector Component
cs'}
  where
    cs :: Vector Component
cs = MixtureModel -> Vector Component
components MixtureModel
m
    cs' :: Vector Component
cs' = (Component -> Component) -> Vector Component -> Vector Component
forall a b. (a -> b) -> Vector a -> Vector b
V.map (String -> Component -> Component
appendNameComponent String
n) Vector Component
cs