Coverage for /home/deng/Projects/metatree_drawer/treeprofiler_algo/pastml/pastml/models/generator.py: 28%

18 statements  

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

1import numpy as np 

2 

3 

4def save_matrix(states, matrix, outfile): 

5 """ 

6 Saves a matrix into a file. 

7 

8 :param states: names of the columns/rows, will be saved in a commented (starting with #) header line before the matrix 

9 :param matrix: np.array matrix to be saved 

10 :param outfile: path to the file where the matrix should be saved 

11 :return: void 

12 """ 

13 np.savetxt(outfile, matrix, delimiter=' ', fmt='%.18e', header=' '.join(states)) 

14 

15 

16def get_diagonalisation(frequencies, rate_matrix=None): 

17 """ 

18 Normalises and diagonalises the rate matrix. 

19 

20 :param frequencies: character state frequencies. 

21 :type frequencies: numpy.array 

22 :param rate_matrix: (optional) rate matrix (by default an all-equal-rate matrix is used) 

23 :type rate_matrix: numpy.ndarray 

24 :return: matrix diagonalisation (d, A, A^{-1}) 

25 such that A.dot(np.diag(d))).dot(A^{-1}) = 1/mu Q (normalised generator) 

26 :rtype: tuple 

27 """ 

28 Q = get_normalised_generator(frequencies, rate_matrix) 

29 d, A = np.linalg.eig(Q) 

30 return d, A, np.linalg.inv(A) 

31 

32 

33def get_normalised_generator(frequencies, rate_matrix=None): 

34 """ 

35 Calculates the normalised generator from the rate matrix and character state frequencies. 

36 

37 :param frequencies: character state frequencies. 

38 :type frequencies: numpy.array 

39 :param rate_matrix: (optional) rate matrix (by default an all-equal-rate matrix is used) 

40 :type rate_matrix: numpy.ndarray 

41 :return: normalised generator 1/mu Q 

42 :rtype: numpy.ndarray 

43 """ 

44 if rate_matrix is None: 

45 n = len(frequencies) 

46 rate_matrix = np.ones(shape=(n, n), dtype=np.float64) - np.eye(n) 

47 generator = rate_matrix * frequencies 

48 generator -= np.diag(generator.sum(axis=1)) 

49 mu = -generator.diagonal().dot(frequencies) 

50 generator /= mu 

51 return generator 

52 

53 

54def get_pij_matrix(t, diag, A, A_inv): 

55 """ 

56 Calculates the probability matrix of substitutions i->j over time t, 

57 given the normalised generator diagonalisation. 

58 

59 

60 :param t: time 

61 :type t: float 

62 :return: probability matrix 

63 :rtype: numpy.ndarray 

64 """ 

65 return A.dot(np.diag(np.exp(diag * t))).dot(A_inv) 

66