spectrum.py 文件源码

def pade(time,dipole):
    damp_const = 100.0
    dipole = np.asarray(dipole) - dipole[0]

    stepsize = time[1] - time[0]
    #print dipole
    damp = np.exp(-(stepsize*np.arange(len(dipole)))/float(damp_const))
    dipole *= damp
    M = len(dipole)
    N = int(np.floor(M / 2))

    #print "N = ", N
    num_pts = 20000
    if N > num_pts:
        N = num_pts
    #print "Trimmed points to: ", N

    # G and d are (N-1) x (N-1)
    # d[k] = -dipole[N+k] for k in range(1,N)
    d = -dipole[N+1:2*N]

    try:
        from scipy.linalg import toeplitz, solve_toeplitz
    except ImportError:
        print("You'll need SciPy version >= 0.17.0")

    try:
        # Instead, form G = (c,r) as toeplitz
        #c = dipole[N:2*N-1]
        #r = np.hstack((dipole[1],dipole[N-1:1:-1]))
        b = solve_toeplitz((dipole[N:2*N-1],\
            np.hstack((dipole[1],dipole[N-1:1:-1]))),d,check_finite=False)
    except np.linalg.linalg.LinAlgError:  
        # OLD CODE: sometimes more stable
        # G[k,m] = dipole[N - m + k] for m,k in range(1,N)
        G = dipole[N + np.arange(1,N)[:,None] - np.arange(1,N)]
        b = np.linalg.solve(G,d)

    # Now make b Nx1 where b0 = 1
    b = np.hstack((1,b))

    # b[m]*dipole[k-m] for k in range(0,N), for m in range(k)
    a = np.dot(np.tril(toeplitz(dipole[0:N])),b)
    p = np.poly1d(a)
    q = np.poly1d(b)

    # If you want energies greater than 2*27.2114 eV, you'll need to change
    # the default frequency range to something greater.

    #frequency = np.arange(0.00,2.0,0.00005)
    frequency = np.arange(0.3,0.75,0.0002)

    W = np.exp(-1j*frequency*stepsize)

    fw = p(W)/q(W)

    return fw, frequency