trajectory.py 文件源码

def _power_series_flat_out(z, t, n, param, y, bound_cond_type):
    """ Provide the power series approximation for x(z,t) and x'(z,t).
    :param z: [0, ..., l] (numpy array)
    :param t: [0, ... , t_end] (numpy array)
    :param n: series termination index (integer)
    :param param: [a2, a1, a0, alpha, beta] (list)
    :param y: flat output with derivation np.array([[y],...,[y^(n/2)]])
    :return: field variable x(z,t) and spatial derivation x'(z,t)
    """
    # TODO: documentation
    a2, a1, a0, alpha, beta = param
    shape = (len(t), len(z))
    x = np.nan*np.ones(shape)
    d_x = np.nan*np.ones(shape)

    # Actually power_series() is designed for robin boundary condition by z=0.
    # With the following modification it can also used for dirichlet boundary condition by z=0.
    if bound_cond_type is 'robin':
        is_robin = 1.
    elif bound_cond_type is 'dirichlet':
        alpha = 1.
        is_robin = 0.
    else:
        raise ValueError("Selected Boundary condition {0} not supported! Use 'robin' or 'dirichlet'".format(
            bound_cond_type))

    # TODO: flip iteration order: z <--> t, result: one or two instead len(t) call's
    for i in range(len(t)):
        sum_x = np.zeros(len(z))
        for j in range(n):
            sum_b = np.zeros(len(z))
            for k in range(j+1):
                sum_b += sm.comb(j, k)*(-a0)**(j-k)*y[k, i]
            sum_x += (is_robin+alpha*z/(2.*j+1.))*z**(2*j)/sm.factorial(2*j)/a2**j*sum_b
        x[i, :] = sum_x

    for i in range(len(t)):
        sum_x = np.zeros(len(z))
        for j in range(n):
            sum_b = np.zeros(len(z))
            for k in range(j+2):
                sum_b += sm.comb(j+1, k)*(-a0)**(j-k+1)*y[k, i]
            if j == 0:
                sum_x += alpha*y[0, i]
            sum_x += (is_robin+alpha*z/(2.*(j+1)))*z**(2*j+1)/sm.factorial(2*j+1)/a2**(j+1)*sum_b
        d_x[i, :] = sum_x

    return x, d_x