def Ct(vZ, dt, index=None, vX=None, vY=None, sci=None, sc=None, avgs=None, mode='full'): """ Calculates the correlation function of a motion given the vX and vZ vectors, in addition to the "avgs" dict, which contains components of the averaged elements of D2 of the inner (bond) frame, evaluated in the current frame. Default mode is 'full', which yields the correlation function considering the residual tensor from a prior motion (tensor normalized by the 0 component). This correlation function is C(t)*A_0 = <D^2_00>A_0+Re(<D^2_10>+<D^2_-10>)Re(A_1)-Im(<D^2_10>+<D^2_-10>)Im(A_1) where the A_p are the averaged components from a prior motion If the mode is set to "separate" or just "s", then the components of the above correlation are returned in a 5xNresxNt array (in the Ct field of the output dictionary) <D^2_00>, to be multiplied by A_0 Re(<D^2_10>-<D^2_-10>), to be multiplied by Re(A_1) -Im(<D^2_10>+<D^2_-10>), to be multiplied by Im(A_1) Re(<D^2_20>+<D^2_-20>), to be multiplied by Re(A_2) Re(<D^2_20>+<D^2_-20>), to be multiplied by Im(A_2) If 'avgs' is provided in this case, then the terms will be multiplied by the appropriate element from 'avgs'. Otherwise, the component is returned without this multiplication. Finally, one may select simply 'P2' (for the second Legendre polynomial), in which case <D^2_00> is returned. If 'avgs' is not given, and mode is not set to 's', then this will be the default option out = Ct(vZ,index,dt,vX=None,vY=None,avgs=None,mode='full') out contains keys 't', 'Ct', 'index', and 'N' """ if index is None: index = np.arange(vZ.shape[1], dtype=int) n = np.size(index) if mode[0].lower() == 's': if avgs is None: rho = np.array([1 + 1j, 1 + 1j, 1, 1 + 1j, 1 + 1j]) else: rho = avgs['rho'] / avgs['rho'][2].real c = [ np.zeros([np.max(index) + 1, np.shape(vZ[0])[1]]) for k in range(5) ] for k in range(n): # scik=sci[:,k] # vXk=vft.R(vX[:,k],*scik) # vYk=vft.R(vY[:,k],*scik) # vZk=vft.R(vZ[:,k],*scik) # vZk_=vft.R(vZ[:,k:],*scik) # CaSb=vXk[0]*vZk_[0]+vXk[1]*vZk_[1]+vXk[2]*vZk_[2] # SaSb=vYk[0]*vZk_[0]+vYk[1]*vZk_[1]+vYk[2]*vZk_[2] # Cb=vZk[0]*vZk_[0]+vZk[1]*vZk_[1]+vZk[2]*vZk_[2] if sc is not None: vXk = vft.R(vX[:, k], *sc[:, k]) vYk = vft.R(vY[:, k], *sc[:, k]) vZk_ = vft.R(vZ[:, k], *sc[:, k:]) vZk = vZk_[:, 0] CaSb = vXk[0] * vZk_[0] + vXk[1] * vZk_[1] + vXk[2] * vZk_[2] SaSb = vYk[0] * vZk_[0] + vYk[1] * vZk_[1] + vYk[2] * vZk_[2] Cb = vZk[0] * vZk_[0] + vZk[1] * vZk_[1] + vZk[2] * vZk_[2] else: CaSb = vX[0, k] * vZ[0, k:] + vX[1, k] * vZ[1, k:] + vX[ 2, k] * vZ[2, k:] SaSb = vY[0, k] * vZ[0, k:] + vY[1, k] * vZ[1, k:] + vY[ 2, k] * vZ[2, k:] Cb = vZ[0, k] * vZ[0, k:] + vZ[1, k] * vZ[1, k:] + vZ[ 2, k] * vZ[2, k:] c[0][index[k:] - index[k]] += (-1 / 2 + 3 / 2 * Cb**2) * rho[2].real c[1][index[k:] - index[k]] += np.sqrt(3 / 2) * (CaSb * Cb) * rho[3].real c[2][index[k:] - index[k]] += np.sqrt(3 / 2) * (SaSb * Cb) * rho[3].imag c[3][index[k:] - index[k]] += np.sqrt( 3 / 8) * (CaSb**2 - SaSb**2) * rho[4].real c[4][index[k:] - index[k]] += np.sqrt(3 / 2) * (CaSb * SaSb) * rho[4].imag if k % np.ceil(n / 100).astype(int) == 0 or k + 1 == n: printProgressBar(k + 1, len(index), prefix='Calculating C(t):', suffix='Complete', length=50) N = get_count(index) #Number of averages for each time point i = N != 0 #Non-zero number of averages N = N[i] #Remove all zeros ct0 = np.array([c0[i].T / N for c0 in c]) if avgs is not None: ct = ct0.sum(axis=0) else: ct = None else: c = np.zeros([np.max(index) + 1, np.shape(vZ[0])[1]]) if avgs is None or mode[:2].lower() == 'p2': for k in range(n): Cb2 = (vZ[0][k:] * vZ[0][k] + vZ[1][k:] * vZ[1][k] + vZ[2][k:] * vZ[2][k])**2 #Cosine beta^2 c[index[k:] - index[k]] += -1 / 2 + 3 / 2 * Cb2 else: if vX is None or vY is None: print('vX and vY required for "full" mode') return rho = avgs['rho'] / avgs['rho'][2].real for k in range(n): # scik=sci[:,k] # vXk=vft.R(vX[:,k],*scik) # vYk=vft.R(vY[:,k],*scik) # vZk=vft.R(vZ[:,k],*scik) # vZk_=vft.R(vZ[:,k:],*scik) # CaSb=vXk[0]*vZk_[0]+vXk[1]*vZk_[1]+vXk[2]*vZk_[2] # SaSb=vYk[0]*vZk_[0]+vYk[1]*vZk_[1]+vYk[2]*vZk_[2] # Cb=vZk[0]*vZk_[0]+vZk[1]*vZk_[1]+vZk[2]*vZk_[2] if sc is not None: vXk = vft.R(vX[:, k], *sc[:, k]) vYk = vft.R(vY[:, k], *sc[:, k]) vZk_ = vft.R(vZ[:, k], *sc[:, k:]) vZk = vZk_[:, 0] CaSb = vXk[0] * vZk_[0] + vXk[1] * vZk_[1] + vXk[2] * vZk_[ 2] SaSb = vYk[0] * vZk_[0] + vYk[1] * vZk_[1] + vYk[2] * vZk_[ 2] Cb = vZk[0] * vZk_[0] + vZk[1] * vZk_[1] + vZk[2] * vZk_[2] else: CaSb = vX[0, k] * vZ[0, k:] + vX[1, k] * vZ[1, k:] + vX[ 2, k] * vZ[2, k:] SaSb = vY[0, k] * vZ[0, k:] + vY[1, k] * vZ[1, k:] + vY[ 2, k] * vZ[2, k:] Cb = vZ[0, k] * vZ[0, k:] + vZ[1, k] * vZ[1, k:] + vZ[ 2, k] * vZ[2, k:] c[index[k:]-index[k]]+=(-1/2+3/2*Cb**2)*rho[2].real+\ np.sqrt(3/2)*(CaSb*Cb)*rho[3].real+\ np.sqrt(3/2)*(SaSb*Cb)*rho[3].imag+\ np.sqrt(3/8)*(CaSb**2-SaSb**2)*rho[4].real+\ np.sqrt(3/2)*(CaSb*SaSb)*rho[4].imag if k % np.ceil(n / 100).astype(int) == 0 or k + 1 == n: printProgressBar(k + 1, len(index), prefix='Calculating C(t)', suffix='Complete', length=50) N = get_count(index) #Number of averages for each time point i = N != 0 #Non-zero number of averages N = N[i] #Remove all zeros ct = c[i].T / N #Remove zeros, normalize ct0 = None t = np.linspace(0, dt * np.max(index), index[-1] + 1) t = t[i] Ct = {'t': t, 'N': N, 'index': index} if ct is not None: Ct['Ct'] = ct if ct0 is not None: Ct['<D2>'] = ct0 return Ct
def ini_vec_load(traj, frame_funs, frame_index=None, index=None, dt=None): """ Loads vectors corresponding to each frame, defined in a list of frame functions. Each element of frame_funs should be a function, which returns one or two vectors. traj should be the trajectory iterable index (optional) is used for sparse sampling of the trajectory dt gives the step size of the MD trajectory (for correcting incorrect step sizes in trajectory) """ if hasattr(frame_funs, '__call__'): frame_funs = [frame_funs ] #In case only one frame defined (unusual usage) nf = len(frame_funs) nt = traj.n_frames if index is None: index = np.arange(nt) if dt is None: dt = traj.dt t = index * dt v = [list() for _ in range(nf)] """v is a list of lists. The outer list runs over the number of frames (length of frame_funs) The inner list runs over the timesteps of the trajectory (that is, the timesteps in index) The inner list contains the results of executing the frame function (outer list) at that time point (inner list) """ for c, i in enumerate(index): traj[i] #Go to current frame for k, f in enumerate(frame_funs): v[k].append(f()) "Print the progress" try: if c % int(len(index) / 100) == 0 or c + 1 == len(index): printProgressBar(c + 1, len(index), prefix='Loading Ref. Frames:', suffix='Complete', length=50) except: pass SZ = list() for k, v0 in enumerate(v): v[k] = np.array(v0) """Put the vectors in order such that if two vectors given, each vector is an element of the first dimension, and the second dimension is X,Y,Z (If only one vector, X,Y,Z is the first dimension) """ if v[k].ndim == 4: v[k] = ((v[k].swapaxes(0, 1)).swapaxes(1, 2)).swapaxes(2, 3) else: v[k] = (v[k].swapaxes(0, 1)).swapaxes(1, 2) SZ.append(v[k].shape[-2]) SZ = np.array(SZ) SZ = SZ[SZ != 1] "Below line causing problems- comment for the moment...means frame_index must be given elsewhere..." # if np.all(SZ==SZ[0]): # frame_index=np.repeat([np.arange(SZ[0])],nf,axis=0) "Somehow, the above line is required for proper iRED functioning...very strange" "!!!! Fix and understand above line's glitch!!!!" return { 'n_frames': nf, 'v': v, 't': t, 'index': index, 'frame_index': frame_index }