wendy.py 文件源码

def _setup_arrays(x,v,m,omega=None):
    sindx= numpy.argsort(x)
    # Keep track of amount of mass above and below and compute acceleration
    mass_below= numpy.cumsum(m[sindx])
    mass_below[-1]= 0.
    mass_below= numpy.roll(mass_below,1)
    mass_above= numpy.cumsum(m[sindx][::-1])[::-1]
    mass_above[0]= 0.
    mass_above= numpy.roll(mass_above,-1)
    a= (mass_above-mass_below)[numpy.argsort(sindx)]
    # Solve for all collisions, using C code for consistency
    tcoll= []
    for xi,vi,ai,xii,vii,aii in zip(x[sindx][:-1],v[sindx][:-1],a[sindx][:-1],
                                    x[sindx][1:],v[sindx][1:],a[sindx][1:]):
        if omega is None:
            tcoll.append(_wendy_solve_coll_quad_func(xi-xii,vi-vii,(ai-aii)/2.))
        else:
            tcoll.append(_wendy_solve_coll_harm_func(xi-xii,vi-vii,ai-aii,omega))
    tcoll= numpy.array(tcoll)
    cindx= ctypes.c_int(numpy.argmin(tcoll))
    next_tcoll= ctypes.c_double(tcoll[cindx])
    # Prepare for C
    err= ctypes.c_int(0)
    #Array requirements
    x= numpy.require(x,dtype=numpy.float64,requirements=['C','W'])
    v= numpy.require(v,dtype=numpy.float64,requirements=['C','W'])
    a= numpy.require(a,dtype=numpy.float64,requirements=['C','W'])
    m= numpy.require(m,dtype=numpy.float64,requirements=['C','W'])
    sindx= numpy.require(sindx,dtype=numpy.int32,requirements=['C','W'])
    tcoll= numpy.require(tcoll,dtype=numpy.float64,requirements=['C','W'])
    return (x,v,m,a,sindx,cindx,next_tcoll,tcoll,err)