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
}