diff --git a/mastermsm/msm/msm.py b/mastermsm/msm/msm.py index e8bf21e..74ffd3c 100644 --- a/mastermsm/msm/msm.py +++ b/mastermsm/msm/msm.py @@ -7,6 +7,7 @@ from functools import reduce, cmp_to_key import itertools import numpy as np +from numpy.core.numeric import normalize_axis_tuple import networkx as nx import scipy.linalg as spla import matplotlib.pyplot as plt @@ -350,7 +351,7 @@ def do_count(self, sliding=True): self.count += self.count.transpose() self.keep_states, self.keep_keys = self.check_connect() - def do_trans(self, evecs=False): + def do_trans(self, evecs=False, normalize=False): """ Wrapper for transition matrix calculation. Also calculates its eigenvalues and right eigenvectors. @@ -365,7 +366,9 @@ def do_trans(self, evecs=False): nkeep = len(self.keep_states) keep_states = self.keep_states count = self.count - self.trans = msm_lib.calc_trans(nkeep, keep_states, count) + ###print('ionix keep_states:',keep_states) + ###print('ionix count matrix:',count) + self.trans = msm_lib.calc_trans(nkeep, keep_states, count, normalize=normalize) if not evecs: self.tauT, self.peqT = self.calc_eigsT() else: diff --git a/mastermsm/msm/msm_lib.py b/mastermsm/msm/msm_lib.py index f5e66c8..928e171 100644 --- a/mastermsm/msm/msm_lib.py +++ b/mastermsm/msm/msm_lib.py @@ -2,6 +2,7 @@ This file is part of the MasterMSM package. """ +import sys #ionix import copy import numpy as np import networkx as nx @@ -536,15 +537,15 @@ def do_boots_worker(x): peqT = rvecsT[:,ieqT]/peqT_sum return tauT, peqT, trans, keep_keys -def calc_trans(nkeep=None, keep_states=None, count=None): +def calc_trans(nkeep=None, keep_states=None, count=None, normalize=False): """ Calculates transition matrix. Uses the maximum likelihood expression by Prinz et al.[1]_ Parameters ---------- - lagt : float - Lag time for construction of MSM. + normalize : boolean + Use method of Zimmerman et al. JCTC 2018 Returns ------- @@ -558,11 +559,23 @@ def calc_trans(nkeep=None, keep_states=None, count=None): Generation and validation", J. Chem. Phys. (2011). """ trans = np.zeros([nkeep, nkeep], float) - for i in range(nkeep): - ni = reduce(lambda x, y: x + y, map(lambda x: - count[keep_states[x]][keep_states[i]], range(nkeep))) - for j in range(nkeep): - trans[j][i] = float(count[keep_states[j]][keep_states[i]])/float(ni) + if normalize: + pseudo = 1./float(nkeep) + for i in range(nkeep): + ni = reduce(lambda x, y: x + y, map(lambda x: + count[keep_states[x]][keep_states[i]], range(nkeep))) + for j in range(nkeep): + trans[j][i] = float(count[keep_states[j]][keep_states[i]]) + pseudo + trans[j][i] /= float(ni + pseudo) + else: + for i in range(nkeep): + ni = reduce(lambda x, y: x + y, map(lambda x: + count[x][i], range(nkeep))) + #ionix count[keep_states[x]][keep_states[i]], range(nkeep))) + if ni==0.0: print(count) #ionix + sys.stdout.flush() + for j in range(nkeep): + trans[j][i] = float(count[keep_states[j]][keep_states[i]])/float(ni) return trans def calc_rate(nkeep, trans, lagt): diff --git a/mastermsm/trajectory/traj.py b/mastermsm/trajectory/traj.py index 27c05da..76a8534 100644 --- a/mastermsm/trajectory/traj.py +++ b/mastermsm/trajectory/traj.py @@ -257,6 +257,35 @@ def discretize(self, method="rama", states=None, nbins=20,\ elif method == "contacts_hdb": self.distraj = traj_lib.discrete_contacts_hdbscan(mcs, ms, self.mdt) + def tica(self, method='contacts', lagt=None): + """ TICA dimensionality reduction. + + Parameters + ---------- + method : str + Compute input for TICA + lagt : int + Lag time for TICA + + Returns + ------- + evals : array + Eigenvalues of TICA coordinate transformation + evecs : array + Eigenvectors of TICA coordinate transformation + + """ + if method == "contacts": + dists = md.compute_contacts(self.mdt, contacts='all', scheme='ca', periodic=True) + x = [] + for dist in dists[0]: + x.append(dist) + x = np.transpose(x) + if lagt == None: lagt = self.dt + evals, evecs = traj_lib.tica_worker(x,lagt) + ###self.distraj = PROJECT trajectory onto TICA vectors? + return evals, evecs + def find_keys(self, exclude=['O']): """ Finds out the discrete states in the trajectory diff --git a/mastermsm/trajectory/traj_lib.py b/mastermsm/trajectory/traj_lib.py index 91257e4..0291013 100644 --- a/mastermsm/trajectory/traj_lib.py +++ b/mastermsm/trajectory/traj_lib.py @@ -434,3 +434,98 @@ def _shift(psi, phi): if psi_s[i] > 2: psi_s[i] -= 2*np.pi return phi_s, psi_s + +def tica_worker(x,tau,dim=2): + """ + Calculate TICA components for features trajectory + 'x' with lag time 'tau'. + Schultze et al. JCTC 2021 + + Parameters + ----------- + x : array + Array with features for each frame of the discrete trajectory. + tau : int + Lag time corresponding to the discrete trajectory. + dim : int + Number of TICA dimensions to be computed. + + Returns: + ------- + evals : numpy array + Resulting eigenvalues + evecs : numpy array + Resulting reaction coordinates + + """ + + # x[0] contiene la lista de los valores del + # primer feature para todos los frames, x[1] + # la lista del segundo feature, etc. + print('Lag time for TICA:',tau) + + # compute mean free x + x = meanfree(x) + # estimate covariant symmetrized matrices + cmat, cmat0 = covmatsym(x,tau) + # solve generalized eigenvalue problem + #n = len(x) + evals, evecs = \ + spla.eig(cmat,b=cmat0,left=True,right=False) + #spla.eigh(cmat,b=cmat0,eigvals_only=False,subset_by_index=[n-dim,n-1]) + + return evals, evecs + +def meanfree(x): + """ + Compute mean free coordinates, i.e. + with zero time average. + + """ + for i,xval in enumerate(x): + x[i] = xval - np.mean(xval) + return x + +def covmatsym(x,tau): + """ + Build symmetrized covariance + matrices. + + """ + cmat = np.zeros((len(x),len(x)),float) + for i,xval in enumerate(x): + for j in range(i): + cmat[i][j] = covmat(xval,x[j],tau) + cmat /= float(len(x[0])-tau) + cmat *= 0.5 + + cmat0 = np.zeros((len(x),len(x)),float) + for i,xval in enumerate(x): + for j in range(i): + cmat0[i][j] = covmat0(xval,x[j],tau) + cmat0 /= float(len(x[0])-tau) + cmat0 *= 0.5 + + return cmat, cmat0 + +def covmat(x,y,tau): + """ + Calculate covariance matrices (right). + + """ + if len(x) != len(y): sys.exit('cannot obtain covariance matrices') + aux = 0.0 + for i,xval in enumerate(x): + aux += xval*y[i+tau] + x[i+tau]*y[i] + if i == (len(x)-tau-1): return aux + +def covmat0(x,y,tau): + """ + Calculate covariance matrices (left). + + """ + if len(x) != len(y): sys.exit('cannot obtain covariance matrices') + aux = 0.0 + for i,xval in enumerate(x): + aux += xval*y[i] + x[i+tau]*y[i+tau] + if i == (len(x)-tau-1): return aux