-- |
-- Module      :  ELynx.MarkovProcess.MixtureModel
-- Description :  Mixture models are a set of substitution models with weights
-- Copyright   :  2021 Dominik Schrempf
-- 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.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.Alphabet.Alphabet hiding (all)
import qualified ELynx.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 -> Double
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
$cshowsPrec :: Int -> Component -> ShowS
showsPrec :: Int -> Component -> ShowS
$cshow :: Component -> String
show :: Component -> String
$cshowList :: [Component] -> ShowS
showList :: [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
$creadsPrec :: Int -> ReadS Component
readsPrec :: Int -> ReadS Component
$creadList :: ReadS [Component]
readList :: ReadS [Component]
$creadPrec :: ReadPrec Component
readPrec :: ReadPrec Component
$creadListPrec :: ReadPrec [Component]
readListPrec :: ReadPrec [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
$cshowsPrec :: Int -> MixtureModel -> ShowS
showsPrec :: Int -> MixtureModel -> ShowS
$cshow :: MixtureModel -> String
show :: MixtureModel -> String
$cshowList :: [MixtureModel] -> ShowS
showList :: [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
$creadsPrec :: Int -> ReadS MixtureModel
readsPrec :: Int -> ReadS MixtureModel
$creadList :: ReadS [MixtureModel]
readList :: ReadS [MixtureModel]
$creadPrec :: ReadPrec MixtureModel
readPrec :: ReadPrec MixtureModel
$creadListPrec :: ReadPrec [MixtureModel]
readListPrec :: ReadPrec [MixtureModel]
Read)

-- | Get weights.
getWeights :: MixtureModel -> V.Vector Weight
getWeights :: MixtureModel -> Vector Double
getWeights = (Component -> Double) -> Vector Component -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> Double
weight (Vector Component -> Vector Double)
-> (MixtureModel -> Vector Component)
-> MixtureModel
-> Vector Double
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

normalizeGlobally :: V.Vector Weight -> V.Vector S.SubstitutionModel -> V.Vector S.SubstitutionModel
normalizeGlobally :: Vector Double
-> Vector SubstitutionModel -> Vector SubstitutionModel
normalizeGlobally Vector Double
ws Vector SubstitutionModel
ss = (SubstitutionModel -> SubstitutionModel)
-> Vector SubstitutionModel -> Vector SubstitutionModel
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Double -> SubstitutionModel -> SubstitutionModel
S.scale (Double -> SubstitutionModel -> SubstitutionModel)
-> Double -> SubstitutionModel -> SubstitutionModel
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Fractional a => a -> a
recip Double
c) Vector SubstitutionModel
ss
  where
    cks :: Vector Double
cks = (SubstitutionModel -> Double)
-> Vector SubstitutionModel -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map SubstitutionModel -> Double
S.totalRate Vector SubstitutionModel
ss
    cNoWeights :: Double
cNoWeights = Vector Double -> Double
forall a. Num a => Vector a -> a
V.sum (Vector Double -> Double) -> Vector Double -> Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Vector Double
ws Vector Double
cks
    c :: Double
c = Double
cNoWeights Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Vector Double -> Double
forall a. Num a => Vector a -> a
V.sum Vector Double
ws

-- | Create a mixture model from a list of substitution models.
--
-- If 'S.Normalize' is 'S.DoNormalize', globally normalize the mixture model.
-- Global normalization has no effect if all components are already normalized.
fromSubstitutionModels ::
  S.Name ->
  S.Normalize ->
  V.Vector Weight ->
  V.Vector S.SubstitutionModel ->
  MixtureModel
fromSubstitutionModels :: String
-> Normalize
-> Vector Double
-> Vector SubstitutionModel
-> MixtureModel
fromSubstitutionModels String
n Normalize
nz Vector Double
ws Vector SubstitutionModel
sms
  | Vector Double -> Bool
forall a. Vector a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Vector Double
ws = String -> MixtureModel
forall a. HasCallStack => String -> a
error String
"fromSubstitutionModels: No weights given."
  | Vector Double -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Double
ws Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector SubstitutionModel -> Int
forall a. Vector a -> 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
    smsNormalized :: Vector SubstitutionModel
smsNormalized = case Normalize
nz of
      Normalize
S.DoNormalize -> Vector Double
-> Vector SubstitutionModel -> Vector SubstitutionModel
normalizeGlobally Vector Double
ws Vector SubstitutionModel
sms
      Normalize
S.DoNotNormalize -> Vector SubstitutionModel
sms
    comps :: Vector Component
comps = (Double -> SubstitutionModel -> Component)
-> Vector Double -> Vector SubstitutionModel -> Vector Component
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Double -> SubstitutionModel -> Component
Component Vector Double
ws Vector SubstitutionModel
smsNormalized
    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
-> Normalize
-> Vector Double
-> Vector SubstitutionModel
-> MixtureModel
fromSubstitutionModels String
n Normalize
S.DoNotNormalize Vector Double
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 Double
ws = (Component -> Double) -> Vector Component -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> Double
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 :: Double -> Component -> Component
scaleComponent Double
s Component
c = Component
c {substModel = s'} where s' :: SubstitutionModel
s' = Double -> SubstitutionModel -> SubstitutionModel
S.scale Double
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 :: Double -> MixtureModel -> MixtureModel
scale Double
s MixtureModel
m = MixtureModel
m {components = 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 (Double -> Component -> Component
scaleComponent Double
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 = Double -> MixtureModel -> MixtureModel
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
c) MixtureModel
mm
  where
    c :: Double
c = Vector Double -> Double
forall a. Num a => Vector a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Vector Double -> Double) -> Vector Double -> Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double)
-> Vector Double -> Vector Double -> Vector Double
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) Vector Double
weights Vector Double
scales
    weights :: Vector Double
weights = MixtureModel -> Vector Double
getWeights MixtureModel
mm
    scales :: Vector Double
scales = (SubstitutionModel -> Double)
-> Vector SubstitutionModel -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map SubstitutionModel -> Double
S.totalRate (Vector SubstitutionModel -> Vector Double)
-> Vector SubstitutionModel -> Vector Double
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 = 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 = 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