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