test_transforms.py 文件源码

def test_axis(ST, quad, axis):
    ST = ST(N, quad=quad, plan=True)
    points, weights = ST.points_and_weights(N)
    f_hat = np.random.random(N)

    B = inner_product((ST, 0), (ST, 0))

    c = np.zeros_like(f_hat)
    c = B.solve(f_hat, c)

    # Multidimensional version
    bc = [np.newaxis,]*3
    bc[axis] = slice(None)
    fk = np.broadcast_to(f_hat[bc], (N,)*3).copy()
    ST.plan((N,)*3, axis, np.float, {})
    if hasattr(ST, 'bc'):
        ST.bc.set_tensor_bcs(ST) # To set Dirichlet boundary conditions on multidimensional array
    ck = np.zeros_like(fk)
    ck = B.solve(fk, ck, axis=axis)
    cc = [0,]*3
    cc[axis] = slice(None)
    assert np.allclose(ck[cc], c)

#test_axis(cbases.ShenDirichletBasis, "GC", 1)