-- |
-- Module      :  ELynx.Data.MarkovProcess.MixtureModel
-- Description :  Mixture models are a set of substitution models with weights
-- Copyright   :  (c) Dominik Schrempf 2020
-- 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.List.NonEmpty as N
import Data.Semigroup
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 -> NonEmpty Component
components :: N.NonEmpty 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 -> N.NonEmpty Weight
getWeights :: MixtureModel -> NonEmpty Weight
getWeights = (Component -> Weight) -> NonEmpty Component -> NonEmpty Weight
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map Component -> Weight
weight (NonEmpty Component -> NonEmpty Weight)
-> (MixtureModel -> NonEmpty Component)
-> MixtureModel
-> NonEmpty Weight
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MixtureModel -> NonEmpty Component
components

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

-- | Create a mixture model from a list of substitution models.
fromSubstitutionModels ::
  S.Name -> N.NonEmpty Weight -> N.NonEmpty S.SubstitutionModel -> MixtureModel
fromSubstitutionModels :: String
-> NonEmpty Weight -> NonEmpty SubstitutionModel -> MixtureModel
fromSubstitutionModels String
n NonEmpty Weight
ws NonEmpty SubstitutionModel
sms =
  if [Alphabet] -> Bool
forall a. Eq a => [a] -> Bool
allEqual ([Alphabet] -> Bool) -> [Alphabet] -> Bool
forall a b. (a -> b) -> a -> b
$ NonEmpty Alphabet -> [Alphabet]
forall a. NonEmpty a -> [a]
N.toList NonEmpty Alphabet
alphs
    then String -> Alphabet -> NonEmpty Component -> MixtureModel
MixtureModel String
n (NonEmpty Alphabet -> Alphabet
forall a. NonEmpty a -> a
N.head NonEmpty Alphabet
alphs) NonEmpty Component
comps
    else
      String -> MixtureModel
forall a. HasCallStack => String -> a
error
        String
"fromSubstitutionModels: alphabets of substitution models are not equal."
  where
    comps :: NonEmpty Component
comps = (Weight -> SubstitutionModel -> Component)
-> NonEmpty Weight
-> NonEmpty SubstitutionModel
-> NonEmpty Component
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
N.zipWith Weight -> SubstitutionModel -> Component
Component NonEmpty Weight
ws NonEmpty SubstitutionModel
sms
    alphs :: NonEmpty Alphabet
alphs = (SubstitutionModel -> Alphabet)
-> NonEmpty SubstitutionModel -> NonEmpty Alphabet
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map SubstitutionModel -> Alphabet
S.alphabet NonEmpty SubstitutionModel
sms
    allEqual :: [a] -> Bool
allEqual [] = Bool
True
    allEqual [a]
xs = (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== [a] -> a
forall a. [a] -> a
head [a]
xs) ([a] -> Bool) -> [a] -> Bool
forall a b. (a -> b) -> a -> b
$ [a] -> [a]
forall a. [a] -> [a]
tail [a]
xs

-- | Concatenate mixture models.
concatenate :: S.Name -> N.NonEmpty MixtureModel -> MixtureModel
concatenate :: String -> NonEmpty MixtureModel -> MixtureModel
concatenate String
n NonEmpty MixtureModel
mms = String
-> NonEmpty Weight -> NonEmpty SubstitutionModel -> MixtureModel
fromSubstitutionModels String
n NonEmpty Weight
ws NonEmpty SubstitutionModel
sms
  where
    comps :: NonEmpty Component
comps = NonEmpty (NonEmpty Component) -> NonEmpty Component
forall a. Semigroup a => NonEmpty a -> a
sconcat (NonEmpty (NonEmpty Component) -> NonEmpty Component)
-> NonEmpty (NonEmpty Component) -> NonEmpty Component
forall a b. (a -> b) -> a -> b
$ (MixtureModel -> NonEmpty Component)
-> NonEmpty MixtureModel -> NonEmpty (NonEmpty Component)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map MixtureModel -> NonEmpty Component
components NonEmpty MixtureModel
mms
    ws :: NonEmpty Weight
ws = (Component -> Weight) -> NonEmpty Component -> NonEmpty Weight
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map Component -> Weight
weight NonEmpty Component
comps
    sms :: NonEmpty SubstitutionModel
sms = (Component -> SubstitutionModel)
-> NonEmpty Component -> NonEmpty SubstitutionModel
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map Component -> SubstitutionModel
substModel NonEmpty 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 :: NonEmpty Component
components = NonEmpty Component
cs'}
  where
    cs :: NonEmpty Component
cs = MixtureModel -> NonEmpty Component
components MixtureModel
m
    cs' :: NonEmpty Component
cs' = (Component -> Component)
-> NonEmpty Component -> NonEmpty Component
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map (Weight -> Component -> Component
scaleComponent Weight
s) NonEmpty 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 = NonEmpty Weight -> Weight
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (NonEmpty Weight -> Weight) -> NonEmpty Weight -> Weight
forall a b. (a -> b) -> a -> b
$ (Weight -> Weight -> Weight)
-> NonEmpty Weight -> NonEmpty Weight -> NonEmpty Weight
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
N.zipWith Weight -> Weight -> Weight
forall a. Num a => a -> a -> a
(*) NonEmpty Weight
weights NonEmpty Weight
scales
    weights :: NonEmpty Weight
weights = MixtureModel -> NonEmpty Weight
getWeights MixtureModel
mm
    scales :: NonEmpty Weight
scales = (SubstitutionModel -> Weight)
-> NonEmpty SubstitutionModel -> NonEmpty Weight
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map SubstitutionModel -> Weight
S.totalRate (NonEmpty SubstitutionModel -> NonEmpty Weight)
-> NonEmpty SubstitutionModel -> NonEmpty Weight
forall a b. (a -> b) -> a -> b
$ MixtureModel -> NonEmpty 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 :: NonEmpty Component
components = NonEmpty Component
cs'}
  where
    cs :: NonEmpty Component
cs = MixtureModel -> NonEmpty Component
components MixtureModel
m
    cs' :: NonEmpty Component
cs' = (Component -> Component)
-> NonEmpty Component -> NonEmpty Component
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map (String -> Component -> Component
appendNameComponent String
n) NonEmpty Component
cs