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 Ct_finF(vZ,nuZ_f,nuXZ_f=None,m=0,mp=0,nuZ_F=None,nuXZ_F=None,index=None): """ Calculates the correlation function for the motion of a bond due to motion of a frame of a frame f in frame F (or in the lab frame) """ "Some checks of the input data" if m!=0 and mp!=0: print('m or mp must be 0') return "Apply frame F" vZF,nuZ_fF,nuXZ_fF=applyFrame(vZ,nuZ_f,nuXZ_f,nuZ_F=nuZ_F,nuXZ_F=nuXZ_F) "Apply frame f" # vZf=applyFrame(vZ,nuZ_F=nuZ_f,nuXZ_F=nuXZ_f) "Axes to project a bond in frame f into frame F" sc=vft.getFrame(nuZ_fF,nuXZ_fF) vZf=vft.R(vZF,*vft.pass2act(*sc)) eFf=[vft.R([1,0,0],*vft.pass2act(*sc)),\ vft.R([0,1,0],*vft.pass2act(*sc)),\ vft.R([0,0,1],*vft.pass2act(*sc))] "Axes of the bond vector in frame F" if not(m==0 and mp==0): sc=vft.getFrame(vZF) vXF=vft.R([1,0,0],*sc) vYF=vft.R([0,1,0],*sc) n=vZ.shape[-1] if vZ.ndim==2: SZ=2*n else: SZ=[vZ.shape[1],2*n] ft0=np.zeros(SZ,dtype=complex) if mp==-2 or m==2: for ea,ax,ay in zip(eFf,vXF,vYF): for eb,bx,by in zip(eFf,vXF,vYF): for eag,gz in zip(ea,vZf): for ebd,dz in zip(eb,vZf): ftee=FT(eag*ebd,index) ft0+=np.sqrt(3/8)*FT(ax*bx*gz*dz-ay*by*gz*dz-1j*2*ax*by*gz*dz,index).conj()*ftee elif mp==-1 or m==1: for ea,ax,ay in zip(eFf,vXF,vYF): for eb,bz in zip(eFf,vZF): for eag,gz in zip(ea,vZf): for ebd,dz in zip(eb,vZf): ftee=FT(eag*ebd,index) ft0+=np.sqrt(3/2)*FT(ax*bz*gz*dz+1j*ay*bz*gz*dz,index).conj()*ftee elif mp==0 and m==0: for ea,az in zip(eFf,vZF): for eb,bz in zip(eFf,vZF): for eag,gz in zip(ea,vZf): for ebd,dz in zip(eb,vZf): ftee=FT(eag*ebd,index) ft0+=3/2*(FT(az*bz*gz*dz,index).conj()*ftee) elif mp==1 or m==-1: for ea,ax,ay in zip(eFf,vXF,vYF): for eb,bz in zip(eFf,vZF): for eag,gz in zip(ea,vZf): for ebd,dz in zip(eb,vZf): ftee=FT(eag*ebd,index) ft0+=np.sqrt(3/2)*FT(-ax*bz*gz*dz+1j*ay*bz*gz*dz,index).conj()*ftee elif mp==2 or m==-2: for ea,ax,ay in zip(eFf,vXF,vYF): for eb,bx,by in zip(eFf,vXF,vYF): for eag,gz in zip(ea,vZf): for ebd,dz in zip(eb,vZf): ftee=FT(eag*ebd,index) ft0+=np.sqrt(3/8)*FT(ax*bx*gz*dz-ay*by*gz*dz+1j*2*ax*by*gz*dz,index).conj()*ftee "Use to properly normalize correlation function" if index is not None: N=get_count(index) else: N=np.arange(n,0,-1) i=N!=0 N=N[i] "Truncate function to half length" if vZ.ndim==3: ct=np.fft.ifft(ft0)[:,:n] ct=ct[:,i]/N else: ct=np.fft.ifft(ft0)[:n] ct=ct[i]/N "Subtract away 1/2 for m,mp=0" if mp==0 and m==0: ct=ct.real-0.5 return ct
def CtPAS(vZ,m=0,mp=0,nuZ_F=None,nuXZ_F=None,index=None): """ Calculates the correlation function for a bond motion, optionally in frame F. Note, this is only for use when we are not interetested in motion of the bond in frame F induced by frame F, rather only the bond motion itself, aligned by F. Defintion of F is optional, in which case the resulting correlation function is the lab frame correlation function. mp is the starting component, and m is the final component (D^2_{mp,m}). One or both must be zero. Currently, only Ct via FT is implemented. Sparse sampling may be used, if specified by index """ "Some checks of the input data" if m!=0 and mp!=0: print('m or mp must be 0') return vZ=applyFrame(vZ,nuZ_F=nuZ_F,nuXZ_F=nuXZ_F) if not(m==0 and mp==0): sc=vft.getFrame(vZ) vX=vft.R([1,0,0],*sc) vY=vft.R([0,1,0],*sc) n=vZ.shape[-1] if vZ.ndim==2: SZ=2*n else: SZ=[vZ.shape[1],2*n] if mp==-2 or m==2: ft0=np.zeros(SZ,dtype=complex) for aX,aY,aZ in zip(vX,vY,vZ): for bX,bY,bZ in zip(vX,vY,vZ): ftzz=FT(aZ*bZ,index) ft0+=np.sqrt(3/8)*FT(aX*bX-aY*bY+1j*2*aX*bY,index).conj()*ftzz elif mp==-1 or m==1: ft0=np.zeros(SZ,dtype=complex) for aX,aY,aZ in zip(vX,vY,vZ): for bZ in vZ: ftzz=FT(aZ*bZ,index) ft0+=np.sqrt(3/2)*FT(aX*bZ-1j*aY*bZ,index).conj()*ftzz elif mp==0 and m==0: ft0=np.zeros(SZ) for aZ in vZ: for bZ in vZ: ftzz=FT(aZ*bZ,index) ft0+=3/2*(ftzz.conj()*ftzz).real elif mp==1 or m==-1: ft0=np.zeros(SZ,dtype=complex) for aX,aY,aZ in zip(vX,vY,vZ): for bZ in vZ: ftzz=FT(aZ*bZ,index) ft0+=np.sqrt(3/2)*FT(-aX*bZ-1j*aY*bZ,index).conj()*ftzz elif mp==2 or m==-2: ft0=np.zeros(SZ,dtype=complex) for aX,aY,aZ in zip(vX,vY,vZ): for bX,bY,bZ in zip(vX,vY,vZ): ftzz=FT(aZ*bZ,index) ft0+=np.sqrt(3/8)*FT(aX*bX-aY*bY-1j*2*aX*bY,index).conj()*ftzz "Truncate function to half length" if vZ.ndim==3: ct=np.fft.ifft(ft0)[:,:n] else: ct=np.fft.ifft(ft0)[:n] "Properly normalize correlation function" if index is not None: N=get_count(index) else: N=np.arange(n,0,-1) ct=ct/N "Subtract away 1/2 for m,mp=0" if mp==0 and m==0: ct=ct.real-0.5 return ct
def Ct_D2inf(vZ,vXZ=None,nuZ_f=None,nuXZ_f=None,nuZ_F=None,nuXZ_F=None,mode='both',m=None,mp=0,index=None): """ Calculates the correlation functions and their values at infinite time simultaneously (greatly reducing the total number of calculations) To perform the calculation in reference frame F, provide nuZ_F and optionally nuXZ_F To calculate the effect of the motion of frame f on the correlation function for the bond, provide nuZ_f and optionally nuXZ_f To only return the correlation function, or only return the values at infinite time, set mode to 'Ct' or 'D2inf', respectively. To only calculate a particular term, set m AND mp (only m required for D2inf). Setting m OR mp will automatically set the other term to 0. Default is for mp=0 (starting component), and m is swept from -2 to 2. Currently, m or mp must be zero index can be provided if the trajectory has been sparsely sampled """ if m!=0 and mp!=0: print('m or mp must be 0') return if m==0: m=mp mp=0 mmpswap=True #We'll just fix this with sign changes at the end else: mmpswap=False #Size of the output n=vZ.shape[-1] if vZ.ndim==2: SZ=[1,2*n] else: SZ=[vZ.shape[1],2*n] #Flags for what we calculate if mode[0].lower()=='b' or mode[0].lower()=='c': calc_ct=True else: calc_ct=False m0=True if m is None or np.abs(m)==0 else False m1=True if m is None or np.abs(m)==1 else False m2=True if m is None or np.abs(m)==2 else False #Pre-allocation if calc_ct: ft0=[np.zeros(SZ,dtype=complex) if m0 else None,\ np.zeros(SZ,dtype=complex) if m1 else None,\ np.zeros(SZ,dtype=complex) if m2 else None] #Pre-allocate storage d20=[np.zeros(SZ[0],dtype=complex) if m0 else None,\ np.zeros(SZ[0],dtype=complex) if m1 else None,\ np.zeros(SZ[0],dtype=complex) if m2 else None] #D2inf is practically free, so we always calculate "l0 is a generator that returns a dictionary with all required components" l0=loop_gen(vZ,vXZ,nuZ_f=nuZ_f,nuXZ_f=nuXZ_f,nuZ_F=nuZ_F,nuXZ_F=nuXZ_F,m=m) if nuZ_f is None: for l in l0: zzp=l['az']*l['bz'] zz=zzp.mean(-1) if m1:p1=np.sqrt(3/2)*(l['ax']*l['bz']+1j*l['ay']*l['bz']) if m2:p2=np.sqrt(3/8)*((l['ax']*l['bx']-l['ay']*l['by'])+1j*2*l['ax']*l['by']) "Calculate the correlation functions (inverse transform at end)" if calc_ct:ftzz=FT(zzp,index=index) if m0 and calc_ct:ft0[0]+=3/2*ftzz*ftzz.conj() if m1 and calc_ct:ft0[1]+=ftzz*FT(p1,index).conj() if m2 and calc_ct:ft0[2]+=ftzz*FT(p2,index).conj() "Calculate the final values" if m0:d20[0]+=3/2*(zz**2) if m1:d20[1]+=zz*(p1.mean(-1)) if m2:d20[2]+=zz*(p2.mean(-1)) else: for l in l0: eep=l['eag']*l['ebd'] #This term appears in all calculations ee=eep.mean(-1) gzdz=l['gz']*l['dz'] #Multiplied by all terms below if m0:p0=3/2*l['az']*l['bz']*gzdz #0,0 term if m1:p1=np.sqrt(3/2)*(l['ax']*l['bz']+1j*l['ay']*l['bz'])*gzdz #+1,0 term if m2:p2=np.sqrt(3/8)*((l['ax']*l['bx']-l['ay']*l['by'])+1j*2*l['ax']*l['by'])*gzdz #+2,0 term "Calculate the correlation functions (inverse transform at end)" if calc_ct:ftee=FT(eep,index=index) if m0 and calc_ct:ft0[0]+=ftee*FT(p0,index).conj() if m1 and calc_ct:ft0[1]+=ftee*FT(p1,index).conj() if m2 and calc_ct:ft0[2]+=ftee*FT(p2,index).conj() "Calculate the final values" if m0:d20[0]+=ee*p0.mean(-1) if m1:d20[1]+=ee*p1.mean(-1) if m2:d20[2]+=ee*p2.mean(-1) if calc_ct: N=get_count(index) if index is not None else np.arange(n,0,-1) i=N!=0 N=N[i] ct=[None if ft is None else (np.fft.ifft(ft,axis=-1)[:,:n])[:,i]/N for ft in ft0] if m0: d20[0]+=-1/2 if calc_ct:ct[0]+=-1/2 "Now correct the signs, fill in other values" if mmpswap: d20=[m_mp_swap(d2,0,k,k,0) for k,d2 in enumerate(d20)] if calc_ct:ct=[m_mp_swap(ct0,0,k,k,0) for k,ct0 in enumerate(ct)] if m is None: d2=np.concatenate([m_mp_swap(d20[2],0,2,0,-2),m_mp_swap(d20[1],0,1,0,-1),d20[0],d20[1],d20[2]]) if calc_ct:ct=np.concatenate([m_mp_swap(ct[2],0,2,0,-2),m_mp_swap(ct[1],0,1,0,-1),ct[0],ct[1],ct[2]]) else: d2=m_mp_swap(d20[np.abs(m)],0,np.abs(m),0,m) if calc_ct:ct=m_mp_swap(ct[np.abs(m)],0,np.abs(m),0,m) if vZ.ndim==2: d2=d2.squeeze() if calc_ct:ct=ct.squeeze() if mode[0].lower()=='b': return ct,d2 elif mode[0].lower()=='c': return ct else: return d2