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
« prev ^ index » next coverage.py v7.2.7, created at 2024-03-21 09:19 +0100
1import numpy as np
4def save_matrix(states, matrix, outfile):
5 """
6 Saves a matrix into a file.
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))
16def get_diagonalisation(frequencies, rate_matrix=None):
17 """
18 Normalises and diagonalises the rate matrix.
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)
33def get_normalised_generator(frequencies, rate_matrix=None):
34 """
35 Calculates the normalised generator from the rate matrix and character state frequencies.
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
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.
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)