Coverage for /home/deng/Projects/metatree_drawer/treeprofiler_algo/pastml/pastml/models/F81Model.py: 50%

14 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-03-21 09:19 +0100

1import numpy as np 

2 

3from pastml.models import ModelWithFrequencies 

4 

5F81 = 'F81' 

6 

7 

8class F81Model(ModelWithFrequencies): 

9 

10 def __init__(self, states, forest_stats, sf=None, frequencies=None, tau=0, 

11 frequency_smoothing=False, optimise_tau=False, parameter_file=None, reoptimise=False, **kwargs): 

12 ModelWithFrequencies.__init__(self, states=states, forest_stats=forest_stats, 

13 sf=sf, tau=tau, frequencies=frequencies, 

14 optimise_tau=optimise_tau, frequency_smoothing=frequency_smoothing, 

15 reoptimise=reoptimise, parameter_file=parameter_file, **kwargs) 

16 self.name = F81 

17 

18 def get_mu(self): 

19 """ 

20 Calculates the mutation rate for F81 (and JC that is a simplification of it), 

21 as \mu = 1 / (1 - sum_i \pi_i^2). This way the overall rate of mutation -\mu trace(\Pi Q) is 1. 

22 See [Gascuel "Mathematics of Evolution and Phylogeny" 2005] for further details. 

23 

24 :return: mutation rate \mu = 1 / (1 - sum_i \pi_i^2) 

25 """ 

26 return 1. / (1. - self.frequencies.dot(self.frequencies)) 

27 

28 def get_Pij_t(self, t, *args, **kwargs): 

29 """ 

30 Calculate the probability of substitution i->j over time t, given the mutation rate mu: 

31 For F81 (and JC which is a simpler version of it) 

32 Pij(t) = \pi_j (1 - exp(-mu t)) + exp(-mu t), if i == j, \pi_j (1 - exp(-mu t)), otherwise 

33 [Gascuel "Mathematics of Evolution and Phylogeny" 2005], 

34 where \pi_i is the equillibrium frequency of state i, 

35 and \mu is the mutation rate: \mu = 1 / (1 - sum_i \pi_i^2). 

36 

37 :param t: time t 

38 :type t: float 

39 :return: probability matrix 

40 :rtype: numpy.ndarray 

41 """ 

42 t = self.transform_t(t) 

43 mu = self.get_mu() 

44 # if mu == inf (e.g. just one state) and t == 0, we should prioritise mu 

45 exp_mu_t = 0. if (mu == np.inf) else np.exp(-mu * t) 

46 return (1 - exp_mu_t) * self.frequencies + np.eye(len(self.frequencies)) * exp_mu_t