-- |
-- 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
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]
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
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]
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 Double
getWeights = forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> Double
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 = forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> SubstitutionModel
substModel 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 = forall a b. (a -> b) -> Vector a -> Vector b
V.map (Double -> SubstitutionModel -> SubstitutionModel
S.scale forall a b. (a -> b) -> a -> b
$ forall a. Fractional a => a -> a
recip Double
c) Vector SubstitutionModel
ss
  where
    cks :: Vector Double
cks = forall a b. (a -> b) -> Vector a -> Vector b
V.map SubstitutionModel -> Double
S.totalRate Vector SubstitutionModel
ss
    cNoWeights :: Double
cNoWeights = forall a. Num a => Vector a -> a
V.sum forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith forall a. Num a => a -> a -> a
(*) Vector Double
ws Vector Double
cks
    c :: Double
c = Double
cNoWeights forall a. Fractional a => a -> a -> a
/ 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
  | forall (t :: * -> *) a. Foldable t => t a -> Bool
null Vector Double
ws = forall a. HasCallStack => String -> a
error String
"fromSubstitutionModels: No weights given."
  | forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector Double
ws forall a. Eq a => a -> a -> Bool
/= forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector SubstitutionModel
sms = forall a. HasCallStack => String -> a
error String
"fromSubstitutionModels: Number of weights and substitution models does not match."
  | Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall {a}. Eq a => Vector a -> Bool
allEqual Vector Alphabet
alphs = 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 (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 = 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 = 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
      | forall a. Vector a -> Bool
V.null Vector a
xs = Bool
True
      | Bool
otherwise = forall a. (a -> Bool) -> Vector a -> Bool
V.all (forall a. Eq a => a -> a -> Bool
== 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 = forall a b. (a -> Vector b) -> Vector a -> Vector b
V.concatMap MixtureModel -> Vector Component
components Vector MixtureModel
mms
    ws :: Vector Double
ws = forall a b. (a -> b) -> Vector a -> Vector b
V.map Component -> Double
weight Vector Component
comps
    sms :: Vector SubstitutionModel
sms = 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 :: SubstitutionModel
substModel = SubstitutionModel
s'} where s' :: SubstitutionModel
s' = Double -> SubstitutionModel -> SubstitutionModel
S.scale Double
s 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 :: Vector Component
components = Vector Component
cs'}
  where
    cs :: Vector Component
cs = MixtureModel -> Vector Component
components MixtureModel
m
    cs' :: Vector Component
cs' = 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 forall a. Fractional a => a -> a -> a
/ Double
c) MixtureModel
mm
  where
    c :: Double
c = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith 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 = forall a b. (a -> b) -> Vector a -> Vector b
V.map SubstitutionModel -> Double
S.totalRate 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 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' = forall a b. (a -> b) -> Vector a -> Vector b
V.map (String -> Component -> Component
appendNameComponent String
n) Vector Component
cs