def frames2data(mol,frame_funs=None,frame_index=None,ffavg='direction',n=100,nr=10,label=None,avg_tensors=True,**kwargs):
"""
Loads a set of frame functions into several data objects (the number of data
objects produced is equal to the number of frames)
Frames may be manually produced and included in a list (frame_funs), or
alternatively can be loaded into the molecule object itself (mol.new_frame)
"""
vecs,avgs,tensors=get_vecs(mol,frame_funs,frame_index,ffavg,n,nr,**kwargs)
if label is None and frame_funs is None:
label=mol._frame_info['label']
data=list()
for q,(v,a,t) in enumerate(zip(vecs,avgs,tensors)):
data.append(vec2data(v,molecule=mol,**kwargs))
if avg_tensors:
if frame_funs is None:frame_funs=mol._vf
if frame_index is None:frame_index=mol._frame_info['frame_index']
if q>0:
"Extract the frame for the first point in the trajectory (index=np.array([0]))"
v0=ini_vec_load(mol.mda_object.trajectory,[frame_funs[q-1]],[frame_index[q-1]],np.array([0]))['v'][0]
if len(v0)==2:
v1,v2=v0
else:
v1,v2=v0,None
v1,inan=apply_index(v1,frame_index[q-1])
v2,_=apply_index(v2,frame_index[q-1])
sc=vft.getFrame(v1,v2) #Angles to rotate first frame
sc=[sc0.squeeze() for sc0 in sc]
"Apply to D2 evaluated at t=infty (D2inf)"
rho1=vft.pars2Spher(a['delta'],a['eta'],*a['euler'])
rho=vft.Rspher(rho1,*sc)
out=vft.Spher2pars(rho,return_angles=True)
data[-1].vars['D2inf']={'delta':out[0],'eta':out[1],'euler':out[2:],'rho':a['rho']}
"Apply to the averaged tensors"
rho1=vft.pars2Spher(t['delta'],t['eta'],*t['euler'])
rho=vft.Rspher(rho1,*sc)
out=vft.Spher2pars(rho,return_angles=True)
data[-1].vars['avg_tensors']={'delta':out[0],'eta':out[1],'euler':out[2:],'A0':t['rho']}
else:
data[-1].vars['D2inf']=a
data[-1].vars['avg_tensors']=t
if label is not None:
data[-1].label=label
return data
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 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 frames2data(mol,
frame_funs=None,
frame_index=None,
mode='full',
n=100,
nr=10,
label=None,
dt=None):
"""
Loads a set of frame functions into several data objects (the number of data
objects produced is equal to the number of frames)
Frames may be manually produced and included in a list (frame_funs), or
alternatively can be loaded into the molecule object itself (mol.new_frame)
"""
if mode.lower() == 'both':
vecs, vecs_z, avgs, tensors = get_vecs(mol,
frame_funs=frame_funs,
frame_index=frame_index,
mode=mode,
n=n,
nr=nr,
dt=dt)
else:
vecs, avgs, tensors = get_vecs(mol,
frame_funs=frame_funs,
frame_index=frame_index,
mode=mode,
n=n,
nr=nr,
dt=dt)
if label is None and frame_funs is None:
label = mol._frame_info['label']
data = list()
"Here, we get the orientation of all frames at t=0"
if frame_funs is None: frame_funs = mol._vf
if frame_index is None: frame_index = mol._frame_info['frame_index']
out = ini_vec_load(mol.mda_object.trajectory,
frame_funs,
frame_index,
index=np.array([0]))
for q, (v, a, t) in enumerate(zip(vecs, avgs, tensors)):
if mode.lower() == 'both':
data.append([vec2data(v,a,mode='full',molecule=mol),\
vec2data(vecs_z[q],mode='z',molecule=mol)])
else:
data.append(vec2data(v, a, mode=mode, molecule=mol))
if q > 0:
v0 = out['v'][q - 1]
v1, v2 = v0 if len(v0) == 2 else (v0, None)
v1, inan = apply_index(v1, frame_index[q - 1])
v2, _ = apply_index(v2, frame_index[q - 1])
sc1 = vft.getFrame(v1, v2)
sc1 = [sc0.squeeze() for sc0 in sc1]
rho = vft.Rspher(t['rho'], *sc1)
new = vft.Spher2pars(rho, return_angles=True)
avg_tensors = {
'delta': new[0],
'eta': new[1],
'euler': new[2:],
'rho': t['rho']
}
else:
avg_tensors = t
out1 = out.copy()
i = np.concatenate((np.arange(q), [len(vecs) - 1]))
out1['v'] = [out1['v'][k] for k in i]
out1['frame_index'] = [out1['frame_index'][k] for k in i]
vecs1, _, _ = applyFrame(out1)
# sc=[vft.getFrame(v['vZ'],v['vX']) for v in vecs1[:-1]]
sc = [vft.pass2act(*v['sc']) for v in vecs1[:-1]]
sc.append(vft.getFrame(vecs1[-1]['vZ']))
rho = a['rho']
for sc0 in sc[::-1]:
rho = vft.Rspher(rho, *sc0)
new = vft.Spher2pars(rho.squeeze(), return_angles=True)
D2inf = {
'delta': new[0],
'eta': new[1],
'euler': new[2:],
'rho': a['rho']
}
if mode.lower() == 'both':
if label is not None:
data[-1][0].label = label
data[-1][1].label = label
data[-1][0].vars['D2inf'] = D2inf
data[-1][1].vars['D2inf'] = D2inf
data[-1][0].vars['avg_tensors'] = avg_tensors
data[-1][1].vars['avg_tensors'] = avg_tensors
else:
if label is not None:
data[-1].label = label
data[-1].vars['D2inf'] = D2inf
data[-1].vars['avg_tensors'] = avg_tensors
if mode.lower() == 'both':
data = [[d[0] for d in data], [d[1] for d in data]]
return data
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 frames2data(mol,
frame_funs=None,
frame_index=None,
mode='full',
n=100,
nr=10,
label=None,
dt=None):
"""
Loads a set of frame functions into several data objects (the number of data
objects produced is equal to the number of frames)
Frames may be manually produced and included in a list (frame_funs), or
alternatively can be loaded into the molecule object itself (mol.new_frame)
"""
vecs, avgs, tensors = get_vecs(mol,
frame_funs=frame_funs,
frame_index=frame_index,
mode=mode,
n=n,
nr=nr,
dt=dt)
if label is None and frame_funs is None:
label = mol._frame_info['label']
data = list()
"Here, we get the orientation of all frames at t=0"
if frame_funs is None: frame_funs = mol._vf
if frame_index is None: frame_index = mol._frame_info['frame_index']
out = ini_vec_load(mol.mda_object.trajectory, frame_funs, frame_index,
np.array([0]))
for q, (v, a, t) in enumerate(zip(vecs, avgs, tensors)):
data.append(vec2data(v, a, mode=mode, molecule=mol))
if q > 0:
v0 = out['v'][q - 1]
v1, v2 = v0 if len(v0) == 2 else (v0, None)
v1, inan = apply_index(v1, frame_index[q - 1])
v2, _ = apply_index(v2, frame_index[q - 1])
sc1 = vft.getFrame(v1, v2)
sc1 = [sc0.squeeze() for sc0 in sc1]
rho = vft.Rspher(t['rho'], *sc1)
new = vft.Spher2pars(rho, return_angles=True)
data[-1].vars['avg_tensors'] = {
'delta': new[0],
'eta': new[1],
'euler': new[2:],
'rho': rho
}
else:
data[-1].vars['avg_tensors'] = t
out1 = out.copy()
i = np.concatenate((np.arange(q), [len(vecs) - 1]))
out1['v'] = [out1['v'][k] for k in i]
out1['frame_index'] = [out1['frame_index'][k] for k in i]
vecs1, _, _ = applyFrame(out1)
sc = [vft.getFrame(v['vZ'], v['vX']) for v in vecs1[:-1]]
sc.append(vft.getFrame(vecs1[-1]['vZ']))
rho = a['rho']
for sc0 in sc[::-1]:
rho = vft.Rspher(rho, *sc0)
new = vft.Spher2pars(rho.squeeze(), return_angles=True)
data[-1].vars['D2inf'] = {
'delta': new[0],
'eta': new[1],
'euler': new[2:],
'rho': rho
}
# if q>0:
# "Extract the frame for the first point in the trajectory (index=np.array([0]))"
# out=ini_vec_load(mol.mda_object.trajectory,[frame_funs[q-1],frame_funs[-1]],[frame_index[q-1],frame_funs[-1]],np.array([0]))['v']
# v0=out[0]
# if len(v0)==2:
# v1,v2=v0
# else:
# v1,v2=v0,None
# v1,inan=apply_index(v1,frame_index[q-1])
# v2,_=apply_index(v2,frame_index[q-1])
#
# sc=vft.getFrame(v1,v2) #Angles to rotate first frame
# sc=[sc0.squeeze() for sc0 in sc]
#
# v0=out[1]
# if len(v0)==2:
# v1,v2=v0
# else:
# v1,v2=v0,None
#
# sc1=vft.getFrame(v1,v2)
# sc1=[sc0.squeeze() for sc0 in sc1]
#
# "Apply to D2 evaluated at t=infty (D2inf)"
# rho1=vft.pars2Spher(a['delta'],a['eta'],*a['euler'])
# rho=vft.Rspher(vft.Rspher(rho1,*vft.pass2act(*sc)),*sc1)
## rho=vft.Rspher(rho1,*sc1)
#
# out=vft.Spher2pars(rho,return_angles=True)
# data[-1].vars['D2inf']={'delta':out[0],'eta':out[1],'euler':out[2:],'rho':a['rho']}
#
# "Apply to the averaged tensors"
# rho1=vft.pars2Spher(t['delta'],t['eta'],*t['euler'])
# rho=vft.Rspher(rho1,*sc)
#
#
# out=vft.Spher2pars(rho,return_angles=True)
# data[-1].vars['avg_tensors']={'delta':out[0],'eta':out[1],'euler':out[2:],'rho':t['rho']}
# else:
# v0=ini_vec_load(mol.mda_object.trajectory,[frame_funs[-1]],[frame_funs[-1]],np.array([0]))['v'][0]
# if len(v0)==2:
# v1,v2=v0
# else:
# v1,v2=v0,None
#
# sc=vft.getFrame(v1,v2) #Angles to rotate first frame
# sc=[sc0.squeeze() for sc0 in sc]
#
# rho1=vft.pars2Spher(a['delta'],a['eta'],*a['euler'])
# rho=vft.Rspher(rho1,*sc)
#
# out=vft.Spher2pars(rho,return_angles=True)
# data[-1].vars['D2inf']={'delta':out[0],'eta':out[1],'euler':out[2:],'rho':a['rho']}
#
# data[-1].vars['avg_tensors']=t
if label is not None:
data[-1].label = label
return data
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