from collections import defaultdict import numpy as np from . import config as ttconf from .seq_utils import alphabets, profile_maps, alphabet_synonyms from .gtr import GTR class GTR_site_specific(GTR): """ Defines General-Time-Reversible model of character evolution that allows for different models at different sites in the alignment """ def __init__(self, seq_len=1, approximate=True, **kwargs): """constructor for site specfic GTR models Parameters ---------- seq_len : int, optional number of sites, determines dimensions of frequency vectors etc approximate : bool, optional use linear interpolation for exponentiated matrices to speed up calcuations **kwargs Description """ self.seq_len=seq_len self.approximate = approximate super(GTR_site_specific, self).__init__(**kwargs) self.is_site_specific=True @property def Q(self): """function that return the product of the transition matrix and the equilibrium frequencies to obtain the rate matrix of the GTR model """ tmp = np.einsum('ia,ij->ija', self.Pi, self.W) diag_vals = np.sum(tmp, axis=0) for x in range(tmp.shape[-1]): np.fill_diagonal(tmp[:,:,x], -diag_vals[:,x]) return tmp def assign_rates(self, mu=1.0, pi=None, W=None): """ Overwrite the GTR model given the provided data Parameters ---------- mu : float Substitution rate W : nxn matrix Substitution matrix pi : n vector Equilibrium frequencies """ if not np.isscalar(mu) and pi is not None and len(pi.shape)==2: if mu.shape[0]!=pi.shape[1]: raise ValueError("GTR_site_specific: length of rate vector (got {}) and equilibrium frequency vector (got {}) must match!".format(mu.shape[0], pi.shape[1])) n = len(self.alphabet) if np.isscalar(mu): self._mu = mu*np.ones(self.seq_len) else: self._mu = np.copy(mu) self.seq_len = mu.shape[0] if pi is not None and pi.shape[0]==n and len(pi.shape)==2: self.seq_len = pi.shape[1] Pi = np.copy(pi) else: if pi is not None: if len(pi)==n: Pi = np.repeat([pi], self.seq_len, axis=0).T else: raise ValueError("GTR_site_specific: length of equilibrium frequency vector (got {}) does not match alphabet length {}".format(len(pi), n)) else: Pi = np.ones(shape=(n,self.seq_len)) self._Pi = Pi/np.sum(Pi, axis=0) if W is None or W.shape!=(n,n): if (W is not None) and W.shape!=(n,n): raise ValueError("GTR_site_specific: Size of substitution matrix (got {}) does not match alphabet length {}".format(W.shape, n)) W = np.ones((n,n)) np.fill_diagonal(W, 0.0) np.fill_diagonal(W, - W.sum(axis=0)) else: W=0.5*(np.copy(W)+np.copy(W).T) np.fill_diagonal(W,0) average_rate = np.einsum('ia,ij,ja',self.Pi, W, self.Pi)/self.seq_len # average_rate = W.dot(avg_pi).dot(avg_pi) self._W = W/average_rate self._mu *=average_rate self.is_site_specific=True self._eig() self._make_expQt_interpolator() @classmethod def random(cls, L=1, avg_mu=1.0, alphabet='nuc', pi_dirichlet_alpha=1, W_dirichlet_alpha=3.0, mu_gamma_alpha=3.0, rng=None): """ Creates a random GTR model Parameters ---------- L : int, optional number of sites for which to generate a model avg_mu : float Substitution rate alphabet : str Alphabet name (should be standard: 'nuc', 'nuc_gap', 'aa', 'aa_gap') pi_dirichlet_alpha : float, optional parameter of dirichlet distribution W_dirichlet_alpha : float, optional parameter of dirichlet distribution mu_gamma_alpha : float, optional parameter of dirichlet distribution Returns ------- GTR_site_specific model with randomly sampled frequencies """ if rng is None: rng = np.random.default_rng() alphabet=alphabets[alphabet] gtr = cls(alphabet=alphabet, seq_len=L) n = gtr.alphabet.shape[0] # Dirichlet distribution == l_1 normalized vector of samples of the Gamma distribution if pi_dirichlet_alpha: pi = 1.0*rng.gamma(pi_dirichlet_alpha, size=(n,L)) else: pi = np.ones((n,L)) pi /= pi.sum(axis=0) if W_dirichlet_alpha: tmp = 1.0*rng.gamma(W_dirichlet_alpha, size=(n,n)) else: tmp = np.ones((n,n)) tmp = np.tril(tmp,k=-1) W = tmp + tmp.T if mu_gamma_alpha: mu = rng.gamma(mu_gamma_alpha, size=(L,)) else: mu = np.ones(L) gtr.assign_rates(mu=mu, pi=pi, W=W) gtr.mu *= avg_mu/np.mean(gtr.average_rate()) return gtr @classmethod def custom(cls, mu=1.0, pi=None, W=None, **kwargs): """ Create a GTR model by specifying the matrix explicitly Parameters ---------- mu : float Substitution rate W : nxn matrix Substitution matrix pi : n vector Equilibrium frequencies **kwargs: Key word arguments to be passed to the constructor Keyword Args ------------ alphabet : str Specify alphabet when applicable. If the alphabet specification is required, but no alphabet is specified, the nucleotide alphabet will be used as default. """ gtr = cls(**kwargs) gtr.assign_rates(mu=mu, pi=pi, W=W) return gtr @classmethod def infer(cls, sub_ija, T_ia, root_state, pc=1.0, gap_limit=0.01, Nit=30, dp=1e-5, **kwargs): r""" Infer a GTR model by specifying the number of transitions and time spent in each character. The basic equation that is being solved is :math:`n_{ij} = pi_i W_{ij} T_j` where :math:`n_{ij}` are the transitions, :math:`pi_i` are the equilibrium state frequencies, :math:`W_{ij}` is the "substitution attempt matrix", while :math:`T_i` is the time on the tree spent in character state :math:`i`. To regularize the process, we add pseudocounts and also need to account for the fact that the root of the tree is in a particular state. the modified equation is :math:`n_{ij} + pc = pi_i W_{ij} (T_j+pc+root\_state)` Parameters ---------- nija : nxn matrix The number of times a change in character state is observed between state j and i at position a Tia :n vector The time spent in each character state at position a root_state : np.array probability that site a is in state i. pc : float Pseudocounts, this determines the lower cutoff on the rate when no substitutions are observed **kwargs: Key word arguments to be passed Keyword Args ------------ alphabet : str Specify alphabet when applicable. If the alphabet specification is required, but no alphabet is specified, the nucleotide alphabet will be used as default. """ from scipy import linalg as LA gtr = cls(**kwargs) gtr.logger("GTR: model inference ",1) q = len(gtr.alphabet) L = sub_ija.shape[-1] n_iter = 0 n_ija = np.copy(sub_ija) n_ija[range(q),range(q),:] = 0 n_ij = n_ija.sum(axis=-1) m_ia = np.sum(n_ija,axis=1) + root_state + pc n_a = n_ija.sum(axis=1).sum(axis=0) + pc Lambda = np.sum(root_state,axis=0) + q*pc p_ia_old=np.zeros((q,L)) p_ia = np.ones((q,L))/q mu_a = np.ones(L) W_ij = np.ones((q,q)) - np.eye(q) while (LA.norm(p_ia_old-p_ia)>dp) and n_itera', p_ia, W_ij, T_ia)) if n_iter >= Nit: gtr.logger('WARNING: maximum number of iterations has been reached in GTR inference',3, warn=True) if LA.norm(p_ia_old-p_ia) > dp: gtr.logger('the iterative scheme has not converged',3,warn=True) if gtr.gap_index is not None: for p in range(p_ia.shape[-1]): if p_ia[gtr.gap_index,p]2: W = np.copy(self.W[:,:,pi]) np.fill_diagonal(W, 0) elif pi==0: np.fill_diagonal(self.W, 0) W=self.W ev, evec, evec_inv = self._eig_single_site(W,self.Pi[:,pi]) eigvals.append(ev) vec.append(evec) vec_inv.append(evec_inv) self.eigenvals = np.array(eigvals).T self.v = np.swapaxes(vec,0,-1) self.v_inv = np.swapaxes(vec_inv, 0,-1) def _make_expQt_interpolator(self): """Function that evaluates the exponentiated substitution matrix at multiple time points and constructs a linear interpolation object """ self.rate_scale = self.average_rate().mean() t_grid = (1.0/self.rate_scale)*np.concatenate((np.linspace(0,.1,11)[:-1], np.linspace(.1,1,21)[:-1], np.linspace(1,5,21)[:-1], np.linspace(5,10,11))) stacked_expQT = np.stack([self._expQt(t) for t in t_grid], axis=0) from scipy.interpolate import interp1d self.expQt_interpolator = interp1d(t_grid, stacked_expQT, axis=0, assume_sorted=True, copy=False, kind='linear') def _expQt(self, t): """Raw numerical matrix exponentiation using the diagonalized matrix. This is the computational bottleneck in many simulations. Parameters ---------- t : float time Returns ------- np.array stack of matrices for each site """ eLambdaT = np.exp(t*self.mu*self.eigenvals) return np.einsum('jia,ja,kja->ika', self.v, eLambdaT, self.v_inv) def expQt(self, t): if t*self.rate_scale<10 and self.approximate: return self.expQt_interpolator(t) else: return self._expQt(t) def prop_t_compressed(self, seq_pair, multiplicity, t, return_log=False): print("NOT IMPEMENTED") def propagate_profile(self, profile, t, return_log=False): """ Compute the probability of the sequence state of the parent at time (t+t0, backwards), given the sequence state of the child (profile) at time t0. Parameters ---------- profile : numpy.array Sequence profile. Shape = (L, a), where L - sequence length, a - alphabet size. t : double Time to propagate return_log: bool If True, return log-probability Returns ------- ` res : np.array Profile of the sequence after time t in the past. Shape = (L, a), where L - sequence length, a - alphabet size. """ Qt = self.expQt(t) res = np.einsum('ai,ija->aj', profile, Qt) return np.log(np.maximum(ttconf.TINY_NUMBER,res)) if return_log else np.maximum(0,res) def evolve(self, profile, t, return_log=False): """ Compute the probability of the sequence state of the child at time t later, given the parent profile. Parameters ---------- profile : numpy.array Sequence profile. Shape = (L, a), where L - sequence length, a - alphabet size. t : double Time to propagate return_log: bool If True, return log-probability Returns ------- res : np.array Profile of the sequence after time t in the future. Shape = (L, a), where L - sequence length, a - alphabet size. """ Qt = self.expQt(t) res = np.einsum('ai,jia->aj', profile, Qt) return np.log(res) if return_log else res def prob_t(self, seq_p, seq_ch, t, pattern_multiplicity = None, return_log=False, ignore_gaps=True): """ Compute the probability to observe seq_ch (child sequence) after time t starting from seq_p (parent sequence). Parameters ---------- seq_p : character array Parent sequence seq_c : character array Child sequence t : double Time (branch len) separating the profiles. pattern_multiplicity : numpy array If sequences are reduced by combining identical alignment patterns, these multplicities need to be accounted for when counting the number of mutations across a branch. If None, all pattern are assumed to occur exactly once. return_log : bool It True, return log-probability. Returns ------- prob : np.array Resulting probability """ if t<0: logP = -ttconf.BIG_NUMBER else: tmp_eQT = self.expQt(t) bad_indices=(tmp_eQT==0) logQt = np.log(tmp_eQT + ttconf.TINY_NUMBER*(bad_indices)) logQt[np.isnan(logQt) | np.isinf(logQt) | bad_indices] = -ttconf.BIG_NUMBER seq_indices_c = np.zeros(len(seq_ch), dtype=int) seq_indices_p = np.zeros(len(seq_p), dtype=int) for ai, a in enumerate(self.alphabet): seq_indices_p[seq_p==a] = ai seq_indices_c[seq_ch==a] = ai if len(logQt.shape)==2: logP = np.sum(logQt[seq_indices_p, seq_indices_c]*pattern_multiplicity) else: logP = np.sum(logQt[seq_indices_p, seq_indices_c, np.arange(len(seq_ch))]*pattern_multiplicity) return logP if return_log else np.exp(logP) def average_rate(self): if self.Pi.shape[1]>1: return np.einsum('a,ia,ij,ja->a',self.mu, self.Pi, self.W, self.Pi) else: return self.mu*np.einsum('ia,ij,ja->a',self.Pi, self.W, self.Pi)