def remove_frame(molecule, frame, sel, filename, steps=[0, -1, 1]):
"""
Calculates and saves a trajectory for which a given frame is aligned
throughout the trajectory. Allows visualization of motion within the frame,
in the absence of the frame motion itself.
Arguments are the molecule object, frame, which specifies which frame,
stored in the molecule object (molecule._vf) should be used (alternatively,
one may give the function directly), sel, an MDAnalysis atom group that
includes all atoms to be stored, filename, which determines where to store
the trajectory, and steps, which determines the start,stop, and step size
over the trajectory (as indices)
Note, the frame should return only a single frame (defined by one or two
vectors)
"""
uni = molecule.mda_object
if frame is None:
def f():
return np.atleast_2d([0, 0, 1]).T
else:
f = frame if hasattr(frame, '__call__') else molecule._vf[frame]
nv = 2 if len(f()) == 2 else 1
if steps[1] == -1: steps[1] = uni.trajectory.n_frames
index = np.arange(steps[0], steps[1], steps[2])
def pos_fun():
return (sel.positions - sel.positions.mean(axis=0)).T
vec = ini_vec_load(uni.trajectory, [f, pos_fun], index=index)
v1, v2 = vec['v'][0] if nv == 2 else vec['v'][0], None
v1 = apply_index(v1, np.zeros(sel.n_atoms))
if v2 is not None: v2 = apply_index(v1, np.zeros(sel.n_atoms))
v0 = vec['v'][1]
sc = vft.getFrame(v1, v2)
v0 = vft.R(v0, *vft.pass2act(*sc)).T
with mda.Writer(os.path.splitext(filename)[0] + '.pdb', sel.n_atoms) as W:
sel.positions = v0[0]
W.write(sel)
with mda.Writer(filename, sel.n_atoms) as W:
for v in v0:
sel.positions = v
W.write(sel)
def applyFrame(*vecs,nuZ_F=None,nuXZ_F=None):
"""
Applies a frame, F, to a set of vectors, *vecs, by rotating such that the
vector nuZ_F lies along the z-axis, and nuXZ_F lies in the xz-plane. Input
is the vectors (as *vecs, so list separately, don't collect in a list), and
the frame, defined by nuZ_F (a vector on the z-axis of the frame), and
optionally nuXZ_F (a vector on xy-axis of the frame). These must be given
as keyword arguments.
vecs_F = applyFrame(*vecs,nuZ_F=nuZ_F,nuXZ_F=None,frame_index=None)
Note, one may also omit the frame application and just apply a frame index
"""
if nuZ_F is None:
out=vecs
else:
sc=vft.pass2act(*vft.getFrame(nuZ_F,nuXZ_F))
out=[None if v is None else vft.R(v,*sc) for v in vecs]
if len(vecs)==1:
return out[0]
else:
return out
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 applyFrame(vecs, mode='full'):
"""
Calculates vectors, which may then be used to determine the influence of
each frame on the overall correlation function (via detector analysis)
For each frame, we will calculate its trajectory with respect to the previous
frame (for example, usually the first frame should be some overall motion
and the last frame should be the bonds themselves. The overall motion is
calculated w.r.t. the lab frame)
ffavg is "Faster-frame averaging", which determines how we take into account
the influence of outer frames on the inner frames.
ffavg='off' : Neglect influence of outer frames, only calculate behavior
of frames with respect to the previous frame
ffavg='direction': Apply motion of a given frame onto a vector pointing
in the same direction as the tensor obtained by averaging
the tensor in the outer frame
ffavg='full' : Apply motion using direction, and also scale
(currently unavailable...)
"""
v0 = vecs['v']
nf = len(v0)
nt = vecs['t'].size
avgs = list()
fi = vecs['frame_index']
nb = fi[0].size
"Make sure all vectors in v0 have two elements"
v0 = [v if len(v) == 2 else [v, None] for v in v0]
"Start by updating vectors with the frame index"
# v0=[[apply_index(v0[k][0],fi[k]),apply_index(v0[k][1],fi[k])] for k in range(nf)]
inan = list()
v1 = list()
for k in range(nf):
a, b = apply_index(v0[k][0], fi[k])
c, _ = apply_index(v0[k][1], fi[k])
inan.append(b)
v1.append([a, c])
v0 = v1
vec_out = list()
avgs = list()
tensors = list()
"This is simply the averaged tensors– Not used, just for storage and later plotting"
v1, v2 = v0[-1]
sc0 = vft.getFrame(v1, v2)
Davg = vft.D2(*sc0).mean(axis=-1) #Average of tensor components of frame 0
out = vft.Spher2pars(Davg) #Convert into delta,eta,euler angles
tensors.append({
'delta': out[0],
'eta': out[1],
'euler': out[2:]
}) #Returns in a dict
tensors[-1].update({'A0': Davg[2]})
"""These are the expectation values for the correlation functions
of the individual components (m=-2->2) evaluated at infinite time
"""
v1, _ = v0[-1] #This is the bond direction in the lab frame
# Davg=vft.getD2inf(v1)
Davg = vft.D2inf_v2(v1)
out = vft.Spher2pars(Davg)
avgs.append({'delta': out[0], 'eta': out[1], 'euler': out[2:]})
avgs[-1].update({'rho': Davg})
"Next sweep from outer frame in"
for k in range(nf - 1):
v1, v2 = v0.pop(
0
) #Here we import the outermost frame (that is, the slowest frame)
i = inan.pop(
0) #Here we remove motion from positions in this frame for NaN
v1[0][i] = 0 #Frames with NaN here point along z for all times
v1[1][i] = 0
v1[2][i] = 1
if v2 is not None: #Frames with NaN here point along x for all times
v2[0][i] = 1
v2[1][i] = 0
v2[2][i] = 0
"Get the euler angles of the outermost frame"
sc = vft.getFrame(
v1, v2
) #Cosines and sines of alpha,beta,gamma (6 values ca,sa,cb,sb,cg,sg)
"""Apply rotations to all inner frames- that is, all inner frames are given in the outer frame
Keep in mind that averaged tensors inside this frame are no longer expressed in the lab frame"
This is later taken care of in
"""
vnew = list()
for v in v0:
v1, v2 = v
"Switch active to passive, apply rotation to all inner vectors"
sci = vft.pass2act(*sc)
vnew.append([vft.R(v1, *sci), vft.R(v2, *sci)])
scx = np.array(vft.getFrame(v0[-1][0])).swapaxes(1, 2)
v0 = vnew #Replace v0 for the next step
v1, _ = v0[-1] #This is the bond direction in the current frame
"Averaged tensors from last frame- Not using this, just for record keeping"
sc0 = vft.getFrame(v1)
Davg = vft.D2(*sc0).mean(
axis=-1) #Average of tensor components of frame 0
out = vft.Spher2pars(Davg) #Convert into delta,eta,euler angles
tensors.append({
'delta': out[0],
'eta': out[1],
'euler': out[2:]
}) #Returns in a dict
tensors[-1].update({'rho': Davg})
"D2 evaluated at infinite time of last frame"
# D2inf=vft.getD2inf(v1) #Average D2 components in current frame at infinite time
D2inf = vft.D2inf_v2(v1)
out = vft.Spher2pars(D2inf)
avgs.append({'delta': out[0], 'eta': out[1], 'euler': out[2:]})
avgs[-1].update({'rho': D2inf})
if mode[0].lower() == 'z':
vZ = vft.R(
np.array([0, 0, 1]),
*avgs[-1]['euler']) #Average of outer frame (usually bond)
vZ = vft.R(np.atleast_3d(vZ).repeat(nt, axis=2),
*sc) #Apply rotation of current frame to outer frame
vZ = vZ.swapaxes(1, 2)
vec_out.append({'vZ': vZ, 't': vecs['t'], 'index': vecs['index']})
else:
# vX=vft.R(np.array([1,0,0]),*sc)
# vY=vft.R(np.array([0,1,0]),*sc)
# vZ=vft.R(np.array([0,0,1]),*sc)
# vX=vX.swapaxes(1,2)
# vY=vY.swapaxes(1,2)
# vZ=vZ.swapaxes(1,2)
#
# sci=np.array(vft.euler_prod([sc,vft.pass2act(*sc0),vft.pass2act(*sc)])).swapaxes(1,2)
#
# vec_out.append({'vZ':vZ,'vX':vX,'vY':vY,'t':vecs['t'],'sci':sci,'index':vecs['index']})
sc = np.array(vft.pass2act(*sc)).swapaxes(1, 2)
# sc0=np.array(sc0).swapaxes(1,2)
vZ = vft.R([0, 0, 1], *scx)
vX = vft.R([1, 0, 0], *scx)
vY = vft.R([0, 1, 0], *scx)
vec_out.append({
'sc': sc,
'vZ': vZ,
'vX': vX,
'vY': vY,
't': vecs['t'],
'index': vecs['index']
})
v, _ = v0.pop(
0
) #Operations on all but bond frame result in normalization, but here we have to enforce length
v = vft.norm(v)
v = v.swapaxes(1, 2)
vec_out.append({'vZ': v, 't': vecs['t'], 'index': vecs['index']})
return vec_out, avgs, tensors
def applyFrame2(vecs,return_avgs=True,tensor_avgs=True):
"""
Calculates vectors, which may then be used to determine the influence of
each frame on the overall correlation function (via detector analysis)
For each frame, we will calculate its trajectory with respect to the previous
frame (for example, usually the first frame should be some overall motion
and the last frame should be the bonds themselves. The overall motion is
calculated w.r.t. the lab frame)
ffavg is "Faster-frame averaging", which determines how we take into account
the influence of outer frames on the inner frames.
ffavg='off' : Neglect influence of outer frames, only calculate behavior
of frames with respect to the previous frame
ffavg='direction': Apply motion of a given frame onto a vector pointing
in the same direction as the tensor obtained by averaging
the tensor in the outer frame
ffavg='full' : Apply motion using direction, and also scale
(currently unavailable...)
"""
v0=vecs['v']
nf=len(v0)
nt=vecs['t'].size
avgs=list()
fi=vecs['frame_index']
nb=fi[0].size
"Make sure all vectors in v0 have two elements"
v0=[v if len(v)==2 else [v,None] for v in v0]
"Start by updating vectors with the frame index"
# v0=[[apply_index(v0[k][0],fi[k]),apply_index(v0[k][1],fi[k])] for k in range(nf)]
inan=list()
v1=list()
for k in range(nf):
a,b=apply_index(v0[k][0],fi[k])
c,_=apply_index(v0[k][1],fi[k])
inan.append(b)
v1.append([a,c])
v0=v1
vec_out=list()
avgs=list()
tensors=list()
if tensor_avgs:
"This is simply the averaged tensors"
v1,v2=v0[-1]
sc0=vft.getFrame(v1,v2)
Davg=vft.D2(*sc0).mean(axis=-1) #Average of tensor components of frame 0
out=vft.Spher2pars(Davg) #Convert into delta,eta,euler angles
tensors.append({'delta':out[0],'eta':out[1],'euler':out[2:]}) #Returns in a dict
tensors[-1].update({'rho':Davg})
if return_avgs:
"""These are the expectation values for the correlation functions
of the individual components (m=-2->2) evaluated at infinite time
"""
v1,_=v0[-1]
D2inf=vft.getD2inf(v1)
sc0=[sc0.squeeze() for sc0 in vft.getFrame(v1[:,:,:1])]
sc0[0]=1 #Set alpha to zero (not -gamma)
sc0[1]=0
A0=D2inf[2]
D2inf=vft.Rspher(D2inf,*sc0)
out=vft.Spher2pars(D2inf)
avgs.append({'delta':out[0],'eta':out[1],'euler':out[2:]})
avgs[-1].update({'rho':D2inf})
"Next sweep from outer frame in"
for k in range(nf-1):
v1,v2=v0.pop(0) #Should have two vectors (or v1 and None)
i=inan.pop(0) #Here we remove motion from positions in this frame for NaN
v1[0][i]=0
v1[1][i]=0
v1[2][i]=1
if v2 is not None:
v2[0][i]=1
v2[1][i]=0
v2[2][i]=0
"Get the euler angles"
sc=vft.getFrame(v1,v2) #Cosines and sines of alpha,beta,gamma (6 values ca,sa,cb,sb,cg,sg)
"Apply rotations to all inner frames"
vnew=list()
for v in v0:
v1,v2=v
"Switch active to passive, apply rotation to all inner vectors"
sci=vft.pass2act(*sc)
vnew.append([vft.R(v1,*sci),vft.R(v2,*sci)])
v0=vnew #Replace v0 for the next step
"Averaged tensor from LAST frame"
if tensor_avgs:
v1,v2=v0[-1]
sc0=vft.getFrame(v1,v2)
Davg=vft.D2(*sc0).mean(axis=-1) #Average of tensor components of frame 0
out=vft.Spher2pars(Davg) #Convert into delta,eta,euler angles
tensors.append({'delta':out[0],'eta':out[1],'euler':out[2:]}) #Returns in a dict
tensors[-1].update({'rho':D2inf})
v1,_=v0[-1]
D2inf=vft.getD2inf(v1)
sc0=[sc0.squeeze() for sc0 in vft.getFrame(v1[:,:,:1])]
sc0[0]=1 #Set alpha to 0 (as opposed to -gamma)
sc0[1]=0
A0=D2inf[2]
D2inf=vft.Rspher(D2inf,*sc0)
out=vft.Spher2pars(D2inf)
avgs.append({'delta':out[0],'eta':out[1],'euler':out[2:]})
avgs[-1].update({'rho':D2inf})
vZ=vft.R(np.array([0,0,1]),*avgs[-1]['euler']) #Average of outer frame (usually bond)
vZ=vft.R(np.atleast_3d(vZ).repeat(nt,axis=2),*sc) #Apply rotation of current frame to outer frame
vX=vft.R(np.array([1,0,0]),*avgs[-1]['euler']) #Average of outer frame (usually bond)
vX=vft.R(np.atleast_3d(vX).repeat(nt,axis=2),*sc) #Apply rotation of current frame to outer frame
vec_out.append({'X':{'X':vX[0].T,'Y':vX[1].T,'Z':vX[2].T},\
'Z':{'X':vZ[0].T,'Y':vZ[1].T,'Z':vZ[2].T},\
't':vecs['t'],'index':vecs['index'],'eta':avgs[-1]['eta']})
"Apply to last frame"
v,_=v0.pop(0)
l=np.sqrt(v[0]**2+v[1]**2+v[2]**2)
v=v/l
vec_out.append({'X':v[0].T,'Y':v[1].T,'Z':v[2].T,'t':vecs['t'],'index':vecs['index']})
if return_avgs and tensor_avgs:
for a in avgs: #Convert to angles
e=a['euler']
alpha,beta,gamma=np.arctan2(e[1],e[0]),np.arctan2(e[3],e[2]),np.arctan2(e[5],e[4])
a['euler']=np.concatenate(([alpha],[beta],[gamma]),axis=0)
for a in tensors: #Convert to angles
e=a['euler']
alpha,beta,gamma=np.arctan2(e[1],e[0]),np.arctan2(e[3],e[2]),np.arctan2(e[5],e[4])
a['euler']=np.concatenate(([alpha],[beta],[gamma]),axis=0)
return vec_out,avgs,tensors
elif tensor_avgs:
for a in tensors: #Convert to angles
e=a['euler']
alpha,beta,gamma=np.arctan2(e[1],e[0]),np.arctan2(e[3],e[2]),np.arctan2(e[5],e[4])
a['euler']=np.concatenate(([alpha],[beta],[gamma]),axis=0)
return vec_out,tensors
elif return_avgs:
for a in avgs: #Convert to angles
e=a['euler']
alpha,beta,gamma=np.arctan2(e[1],e[0]),np.arctan2(e[3],e[2]),np.arctan2(e[5],e[4])
a['euler']=np.concatenate(([alpha],[beta],[gamma]),axis=0)
return vec_out,avgs
else:
return vec_out
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_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 D2infPAS(vZ,nuZ_F=None,nuXZ_F=None,m=None):
"""
Estimates the Wigner rotation matrix elements of a bond reorientation between
two times separated by infinite time, that is, we calculate:
Am=lim_t->infty <D_{0m}^2(\Omega_tau,t+tau>
for a bond with direction given by vZ.
nuZ_F and nuXY_F are optional arguments to define frame in which to calculate
the terms. Otherwise, calculation is performed in the input frame of vZ.
A frame_index must also be provided if nuZ_F is not the same size as vZ
Setting m=None (default) results in all terms being returned in an array
(outer dimension runs from m=-2 to 2)
Am = D2inf(vZ,vX=None,vY=None,m=None)
"""
"Apply the frame if required"
vZ=applyFrame(vZ,nuZ_F=nuZ_F,nuXZ_F=nuXZ_F)
if m is None:
m1=[-2,-1,0] #Default, get all values (we'll calculate 1 and 2 from -1 and -2)
else:
m1=[m]
if m!=0:
sc=vft.getFrame(vZ)
vX=vft.R([1,0,0],*sc)
vY=vft.R([0,1,0],*sc)
A=list() #Collect the output
for m0 in m1: #Sweep over the components (or just one in case m is not None)
if vZ.ndim==2:
d2inf=0
else:
d2inf=np.zeros(vZ.shape[1],dtype=complex)
if m0==-2:
for aX,aY,aZ in zip(vX,vY,vZ): #Loop over x,y,z components of all 3 axes
for bX,bY,bZ in zip(vX,vY,vZ): #ditto
d2inf+=np.sqrt(3/8)*((aX*bX).mean(-1)-(aY*bY).mean(-1))*(aZ*bZ).mean(-1)\
+1j*np.sqrt(3/2)*(aX*bY).mean(-1)*(aZ*bZ).mean(-1)
elif m0==-1:
for aX,aY,aZ in zip(vX,vY,vZ): #Loop over x,y,z components of all 3 axes
for bZ in vZ: #ditto
d2inf+=np.sqrt(3/2)*(aX*bZ).mean(-1)*(aZ*bZ).mean(-1)\
-1j*np.sqrt(3/2)*(aY*aZ).mean(-1)*(aZ*bZ).mean(-1)
elif m0==0:
if vZ.ndim==2:
d2inf=-0.5 #Offset for 0,0 term
else:
d2inf[:]=-0.5
for aZ in vZ: #Loop over x,y,z components of all 3 axes
for bZ in vZ: #ditto
d2inf+=3/2*((aZ*bZ).mean(-1))**2 #Real only
elif m0==1:
for aX,aY,aZ in zip(vX,vY,vZ): #Loop over x,y,z components of all 3 axes
for bZ in vZ: #ditto
d2inf+=-np.sqrt(3/2)*(aX*bZ).mean(-1)*(aZ*bZ).mean(-1)\
-1j*np.sqrt(3/2)*(aY*aZ).mean(-1)*(aZ*bZ).mean(-1)
elif m0==2:
for aX,aY,aZ in zip(vX,vY,vZ): #Loop over x,y,z components of all 3 axes
for bX,bY,bZ in zip(vX,vY,vZ): #ditto
d2inf+=np.sqrt(3/8)*((aX*bX).mean(-1)-(aY*bY).mean(-1))*(aZ*bZ).mean(-1)\
-1j*np.sqrt(3/2)*(aX*bY).mean(-1)*(aZ*bZ).mean(-1)
A.append(d2inf)
if m is None:
A.append(-A[1].conj()) #A_1=-A^*_-1
A.append(A[0].conj()) #A_2=A^*_-2
else:
A=A[0] #Only one element- get rid of list
return np.array(A) #Return as numpy array
def loop_gen(vZ,vXZ=None,nuZ_f=None,nuXZ_f=None,nuZ_F=None,nuXZ_F=None,m=None):
"""
Generates the loop items required for calculating correlation functions and
equilibrium values.
For f in F calculations, vZ and nuZ_f are required, with additional optional
arguments nuXZ_f, nuZ_F (calc in LF if omitted), nuXZ_F. If nuZ_f is included,
a dictionary with the following keys is returned for each of 81 iterator
elements
Dictionary keys:
eag: "gamma" element of the vector representing the "alpha" axis for
projection from frame f into frame F
ebd: "delta" element of the vector representing the "beta" axis for
projection from frame f into frame F
gz: "gamma" element of the bond vector in frame f
dz: "delta" element of the bond vector in frame f
ax: "alpha" element of the x-axis of the bond axis system in frame F
ay: "alpha" element of the y-axis of the bond axis system in frame F
az: "alpha" element of the z-axis of the bond axis system in frame F
bx: "beta" element of the x-axis of the bond axis system in frame F
by: "beta" element of the y-axis of the bond axis system in frame F
bz: "beta" element of the z-axis of the bond axis system in frame F
For PAS calculations (optionally in frame F), vZ is required, with additional
optional arguments nuZ_F (calc in LF if omitted) and nuXZ_F. In this case,
a dictionary is returned with the following keys for each of 9 iterator
elements
Dictionary keys:
ax: "alpha" element of the x-axis of the bond axis system
ay: "alpha" element of the y-axis of the bond axis system
az: "alpha" element of the z-axis of the bond axis system
bx: "beta" element of the x-axis of the bond axis system
by: "beta" element of the y-axis of the bond axis system
bz: "beta" element of the z-axis of the bond axis system
If m is provided, only the dictionary elements required for calculating the mth
component will be returned
loop_gen(vZ,nuZ_f=None,nuXZ_f=None,nuZ_F=None,nuXZ_F=None,m=None)
"""
"Apply frame F"
vZF,vXZF,nuZ_fF,nuXZ_fF=applyFrame(vZ,vXZ,nuZ_f,nuXZ_f,nuZ_F=nuZ_F,nuXZ_F=nuXZ_F)
if m!=0:
sc=vft.getFrame(vZF,vXZF)
vXF=vft.R([1,0,0],*sc) #x and y axes of the bond axes in F
vYF=vft.R([0,1,0],*sc)
if nuZ_f is None:
if m is None or m==-2 or m==2:
for ax,ay,az in zip(vXF,vYF,vZF):
for bx,by,bz in zip(vXF,vYF,vZF):
out={'ax':ax,'ay':ay,'az':az,'bx':bx,'by':by,'bz':bz}
yield out
elif m==-1 or m==1:
for ax,ay,az in zip(vXF,vYF,vZF):
for bz in vZF:
out={'ax':ax,'ay':ay,'az':az,'bz':bz}
yield out
elif m==0:
for az in vZF:
for bz in vZF:
out={'az':az,'bz':bz}
yield out
else:
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))]
if m is None:
for ea,ax,ay,az in zip(eFf,vXF,vYF,vZF):
for eb,bx,by,bz in zip(eFf,vXF,vYF,vZF):
for eag,gz in zip(ea,vZf):
for ebd,dz in zip(eb,vZf):
out={'eag':eag,'ebd':ebd,'ax':ax,'ay':ay,'az':az,\
'bx':bx,'by':by,'bz':bz,'gz':gz,'dz':dz}
yield out
elif m==-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(ea,vZf):
out={'eag':eag,'ebd':ebd,'ax':ax,'ay':ay,\
'bx':bx,'by':by,'gz':gz,'dz':dz}
yield out
elif m==-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):
out={'eag':eag,'ebd':ebd,'ax':ax,'ay':ay,\
'bz':bz,'gz':gz,'dz':dz}
yield out
elif 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):
out={'eag':eag,'ebd':ebd,'az':az,\
'bz':bz,'gz':gz,'dz':dz}
yield out