Ejemplo n.º 1
0
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)
Ejemplo n.º 2
0
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
Ejemplo n.º 3
0
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
Ejemplo n.º 4
0
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
Ejemplo n.º 5
0
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
Ejemplo n.º 6
0
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                    
Ejemplo n.º 7
0
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