-- |
-- Module      :  ELynx.Data.MarkovProcess.GammaRateHeterogeneity
-- Description :  Discrete gamma rate heterogeneity
-- Copyright   :  (c) Dominik Schrempf 2020
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu Feb 28 14:09:11 2019.
--
-- At the moment, a mixture model is used to emulate gamma rate heterogeneity. This
-- does not come with huge run time increases when simulating data. For inference
-- however, it would make a lot of sense to reuse the Eigendecomposition for all
-- rate heterogeneity components though.
module ELynx.Data.MarkovProcess.GammaRateHeterogeneity
  ( summarizeGammaRateHeterogeneity,
    expand,
  )
where

import qualified Data.ByteString.Lazy.Char8 as BL
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as N
import qualified ELynx.Data.MarkovProcess.MixtureModel as M
import qualified ELynx.Data.MarkovProcess.PhyloModel as P
import qualified ELynx.Data.MarkovProcess.SubstitutionModel as S
import Numeric.Integration.TanhSinh
import Statistics.Distribution
import Statistics.Distribution.Gamma
import Prelude hiding (repeat)

-- | Short summary of gamma rate heterogeneity parameters.
summarizeGammaRateHeterogeneity :: Int -> Double -> [BL.ByteString]
summarizeGammaRateHeterogeneity :: Int -> Double -> [ByteString]
summarizeGammaRateHeterogeneity Int
n Double
alpha =
  ([Char] -> ByteString) -> [[Char]] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map
    [Char] -> ByteString
BL.pack
    [ [Char]
"Discrete gamma rate heterogeneity.",
      [Char]
"Number of categories: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n,
      [Char]
"Shape parameter of gamma distribution: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Double -> [Char]
forall a. Show a => a -> [Char]
show Double
alpha,
      [Char]
"Rates: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ NonEmpty Double -> [Char]
forall a. Show a => a -> [Char]
show (Int -> Double -> NonEmpty Double
getMeans Int
n Double
alpha)
    ]

-- | For a given number of rate categories, a gamma shape parameter alpha and a
-- substitution model, compute the scaled substitution models corresponding to
-- the gamma rates.
expand :: Int -> Double -> P.PhyloModel -> P.PhyloModel
expand :: Int -> Double -> PhyloModel -> PhyloModel
expand Int
n Double
alpha (P.SubstitutionModel SubstitutionModel
sm) =
  MixtureModel -> PhyloModel
P.MixtureModel (MixtureModel -> PhyloModel) -> MixtureModel -> PhyloModel
forall a b. (a -> b) -> a -> b
$ Int -> Double -> SubstitutionModel -> MixtureModel
expandSubstitutionModel Int
n Double
alpha SubstitutionModel
sm
expand Int
n Double
alpha (P.MixtureModel MixtureModel
mm) =
  MixtureModel -> PhyloModel
P.MixtureModel (MixtureModel -> PhyloModel) -> MixtureModel -> PhyloModel
forall a b. (a -> b) -> a -> b
$ Int -> Double -> MixtureModel -> MixtureModel
expandMixtureModel Int
n Double
alpha MixtureModel
mm

getName :: Int -> Double -> String
getName :: Int -> Double -> [Char]
getName Int
n Double
alpha =
  [Char]
" with discrete gamma rate heterogeneity; "
    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n
    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" categories; "
    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"shape parameter "
    [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Double -> [Char]
forall a. Show a => a -> [Char]
show Double
alpha

splitSubstitutionModel ::
  Int -> Double -> S.SubstitutionModel -> NonEmpty S.SubstitutionModel
splitSubstitutionModel :: Int -> Double -> SubstitutionModel -> NonEmpty SubstitutionModel
splitSubstitutionModel Int
n Double
alpha SubstitutionModel
sm = NonEmpty SubstitutionModel
renamedSMs
  where
    means :: NonEmpty Double
means = Int -> Double -> NonEmpty Double
getMeans Int
n Double
alpha
    scaledSMs :: NonEmpty SubstitutionModel
scaledSMs = (Double -> SubstitutionModel)
-> NonEmpty Double -> NonEmpty SubstitutionModel
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map (Double -> SubstitutionModel -> SubstitutionModel
`S.scale` SubstitutionModel
sm) NonEmpty Double
means
    names :: NonEmpty [Char]
names = [[Char]] -> NonEmpty [Char]
forall a. [a] -> NonEmpty a
N.fromList ([[Char]] -> NonEmpty [Char]) -> [[Char]] -> NonEmpty [Char]
forall a b. (a -> b) -> a -> b
$ (Int -> [Char]) -> [Int] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (([Char]
"; gamma rate category " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++) ([Char] -> [Char]) -> (Int -> [Char]) -> Int -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Char]
forall a. Show a => a -> [Char]
show) [Int
1 :: Int ..]
    renamedSMs :: NonEmpty SubstitutionModel
renamedSMs = ([Char] -> SubstitutionModel -> SubstitutionModel)
-> NonEmpty [Char]
-> NonEmpty SubstitutionModel
-> NonEmpty SubstitutionModel
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
N.zipWith [Char] -> SubstitutionModel -> SubstitutionModel
S.appendName NonEmpty [Char]
names NonEmpty SubstitutionModel
scaledSMs

expandSubstitutionModel ::
  Int -> Double -> S.SubstitutionModel -> M.MixtureModel
expandSubstitutionModel :: Int -> Double -> SubstitutionModel -> MixtureModel
expandSubstitutionModel Int
n Double
alpha SubstitutionModel
sm = [Char]
-> NonEmpty Double -> NonEmpty SubstitutionModel -> MixtureModel
M.fromSubstitutionModels [Char]
name NonEmpty Double
ws NonEmpty SubstitutionModel
sms
  where
    name :: [Char]
name = SubstitutionModel -> [Char]
S.name SubstitutionModel
sm [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> Double -> [Char]
getName Int
n Double
alpha
    ws :: NonEmpty Double
ws = Double -> NonEmpty Double
forall a. a -> NonEmpty a
N.repeat Double
1.0
    sms :: NonEmpty SubstitutionModel
sms = Int -> Double -> SubstitutionModel -> NonEmpty SubstitutionModel
splitSubstitutionModel Int
n Double
alpha SubstitutionModel
sm

expandMixtureModel :: Int -> Double -> M.MixtureModel -> M.MixtureModel
expandMixtureModel :: Int -> Double -> MixtureModel -> MixtureModel
expandMixtureModel Int
n Double
alpha MixtureModel
mm = [Char] -> NonEmpty MixtureModel -> MixtureModel
M.concatenate [Char]
name NonEmpty MixtureModel
renamedMMs
  where
    name :: [Char]
name = MixtureModel -> [Char]
M.name MixtureModel
mm [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Int -> Double -> [Char]
getName Int
n Double
alpha
    means :: NonEmpty Double
means = Int -> Double -> NonEmpty Double
getMeans Int
n Double
alpha
    scaledMMs :: NonEmpty MixtureModel
scaledMMs = (Double -> MixtureModel)
-> NonEmpty Double -> NonEmpty MixtureModel
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
N.map (Double -> MixtureModel -> MixtureModel
`M.scale` MixtureModel
mm) NonEmpty Double
means
    names :: NonEmpty [Char]
names = [[Char]] -> NonEmpty [Char]
forall a. [a] -> NonEmpty a
N.fromList ([[Char]] -> NonEmpty [Char]) -> [[Char]] -> NonEmpty [Char]
forall a b. (a -> b) -> a -> b
$ (Int -> [Char]) -> [Int] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (([Char]
"; gamma rate category " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++) ([Char] -> [Char]) -> (Int -> [Char]) -> Int -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Char]
forall a. Show a => a -> [Char]
show) [Int
1 :: Int ..]
    renamedMMs :: NonEmpty MixtureModel
renamedMMs = ([Char] -> MixtureModel -> MixtureModel)
-> NonEmpty [Char]
-> NonEmpty MixtureModel
-> NonEmpty MixtureModel
forall a b c.
(a -> b -> c) -> NonEmpty a -> NonEmpty b -> NonEmpty c
N.zipWith [Char] -> MixtureModel -> MixtureModel
M.appendNameComponents NonEmpty [Char]
names NonEmpty MixtureModel
scaledMMs

-- For a given number of rate categories 'n' and a shape parameter 'alpha' (the
-- rate or scale is set such that the mean is 1.0), return a list of rates that
-- represent the respective categories. Use the mean rate for each category.
getMeans :: Int -> Double -> NonEmpty Double
getMeans :: Int -> Double -> NonEmpty Double
getMeans Int
n Double
alpha = [Double] -> NonEmpty Double
forall a. [a] -> NonEmpty a
N.fromList ([Double] -> NonEmpty Double) -> [Double] -> NonEmpty Double
forall a b. (a -> b) -> a -> b
$ [Double]
means [Double] -> [Double] -> [Double]
forall a. Semigroup a => a -> a -> a
<> Double -> [Double]
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
lastMean
  where
    gamma :: GammaDistribution
gamma = Double -> Double -> GammaDistribution
gammaDistr Double
alpha (Double
1.0 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
alpha)
    quantiles :: [Double]
quantiles =
      [GammaDistribution -> Double -> Double
forall d. ContDistr d => d -> Double -> Double
quantile GammaDistribution
gamma (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) | Int
i <- [Int
0 .. Int
n]]
    -- Calculate the mean rate. Multiplication with the number of rate
    -- categories 'n' is necessary because in each n-quantile the
    -- probability mass is 1/n.
    meanFunc :: Double -> Double
meanFunc Double
x = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
x Double -> Double -> Double
forall a. Num a => a -> a -> a
* GammaDistribution -> Double -> Double
forall d. ContDistr d => d -> Double -> Double
density GammaDistribution
gamma Double
x
    -- Only calculate the first (n-1) categories with normal integration.
    means :: [Double]
means =
      [ (Double -> Double) -> Double -> Double -> Double
integralAToB Double -> Double
meanFunc ([Double]
quantiles [Double] -> Int -> Double
forall a. [a] -> Int -> a
!! Int
i) ([Double]
quantiles [Double] -> Int -> Double
forall a. [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
      ]
    -- The last category has to be calculated with an improper integration.
    lastMean :: Double
lastMean = (Double -> Double) -> Double -> Double
integralAToInf Double -> Double
meanFunc ([Double]
quantiles [Double] -> Int -> Double
forall a. [a] -> Int -> a
!! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))

eps :: Double
eps :: Double
eps = Double
1e-8

-- The integration method to use
method :: (Double -> Double) -> Double -> Double -> [Result]
method :: (Double -> Double) -> Double -> Double -> [Result]
method = (Double -> Double) -> Double -> Double -> [Result]
parSimpson

-- Helper function for a normal integral from 'a' to 'b'.
integralAToB :: (Double -> Double) -> Double -> Double -> Double
integralAToB :: (Double -> Double) -> Double -> Double -> Double
integralAToB Double -> Double
f Double
a Double
b = Result -> Double
result (Result -> Double) -> ([Result] -> Result) -> [Result] -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> [Result] -> Result
absolute Double
eps ([Result] -> Double) -> [Result] -> Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double) -> Double -> Double -> [Result]
method Double -> Double
f Double
a Double
b

-- Helper function for an improper integral from 'a' to infinity.
integralAToInf :: (Double -> Double) -> Double -> Double
integralAToInf :: (Double -> Double) -> Double -> Double
integralAToInf Double -> Double
f Double
a =
  (Result -> Double
result (Result -> Double) -> ([Result] -> Result) -> [Result] -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> [Result] -> Result
absolute Double
eps ([Result] -> Double) -> [Result] -> Double
forall a b. (a -> b) -> a -> b
$ ((Double -> Double) -> Double -> Double -> [Result])
-> (Double -> Double) -> [Result]
forall r.
((Double -> Double) -> Double -> Double -> r)
-> (Double -> Double) -> r
nonNegative (Double -> Double) -> Double -> Double -> [Result]
method Double -> Double
f) Double -> Double -> Double
forall a. Num a => a -> a -> a
- (Double -> Double) -> Double -> Double -> Double
integralAToB Double -> Double
f Double
0 Double
a