Example #1
0
def test_project_hermite(typecode, dim, ST, quad):
    # Using sympy to compute an analytical solution
    x, y, z = symbols("x,y,z")
    sizes = (20, 19)

    funcs = {
        (1, 0): (cos(4*y)*sin(2*x))*exp(-x**2/2),
        (1, 1): (cos(4*x)*sin(2*y))*exp(-y**2/2),
        (2, 0): (sin(3*z)*cos(4*y)*sin(2*x))*exp(-x**2/2),
        (2, 1): (sin(2*z)*cos(4*x)*sin(2*y))*exp(-y**2/2),
        (2, 2): (sin(2*x)*cos(4*y)*sin(2*z))*exp(-z**2/2)
        }
    syms = {1: (x, y), 2:(x, y, z)}
    xs = {0:x, 1:y, 2:z}

    for shape in product(*([sizes]*dim)):
        bases = []
        for n in shape[:-1]:
            bases.append(Basis(n, 'F', dtype=typecode.upper()))
        bases.append(Basis(shape[-1], 'F', dtype=typecode))

        for axis in range(dim+1):
            ST0 = ST(3*shape[-1], quad=quad)
            bases.insert(axis, ST0)
            fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
            X = fft.local_mesh(True)
            ue = funcs[(dim, axis)]
            u_h = project(ue, fft)
            ul = lambdify(syms[dim], ue, 'numpy')
            uq = ul(*X).astype(typecode)
            uf = u_h.backward()
            assert np.linalg.norm(uq-uf) < 1e-5
            bases.pop(axis)
            fft.destroy()
def test_project_hermite(typecode, dim):
    # Using sympy to compute an analytical solution
    x, y, z = symbols("x,y,z")
    sizes = (20, 19)

    funcs = {
        (1, 0): (cos(4*y)*sin(2*x))*exp(-x**2/2),
        (1, 1): (cos(4*x)*sin(2*y))*exp(-y**2/2),
        (2, 0): (sin(3*z)*cos(4*y)*sin(2*x))*exp(-x**2/2),
        (2, 1): (sin(2*z)*cos(4*x)*sin(2*y))*exp(-y**2/2),
        (2, 2): (sin(2*x)*cos(4*y)*sin(2*z))*exp(-z**2/2)
        }
    xs = {0:x, 1:y, 2:z}

    for shape in product(*([sizes]*dim)):
        bases = []
        for n in shape[:-1]:
            bases.append(FunctionSpace(n, 'F', dtype=typecode.upper()))
        bases.append(FunctionSpace(shape[-1], 'F', dtype=typecode))

        for axis in range(dim+1):
            ST0 = hBasis[0](3*shape[-1])
            bases.insert(axis, ST0)
            fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
            X = fft.local_mesh(True)
            ue = funcs[(dim, axis)]
            due = ue.diff(xs[0], 1)
            u_h = project(ue, fft)
            du_h = project(due, fft)
            du2 = project(Dx(u_h, 0, 1), fft)

            assert np.linalg.norm(du_h-du2) < 1e-5
            bases.pop(axis)
            fft.destroy()
Example #3
0
def test_assign(fam):
    x, y = symbols("x,y")
    for bc in (None, 'Dirichlet', 'Biharmonic'):
        dtype = 'D' if fam == 'F' else 'd'
        bc = 'periodic' if fam == 'F' else bc
        if bc == 'Biharmonic' and fam in ('La', 'H'):
            continue
        tol = 1e-12 if fam in ('C', 'L', 'F') else 1e-5
        N = (10, 12)
        B0 = Basis(N[0], fam, dtype=dtype, bc=bc)
        B1 = Basis(N[1], fam, dtype=dtype, bc=bc)
        u_hat = Function(B0)
        u_hat[1:4] = 1
        ub_hat = Function(B1)
        u_hat.assign(ub_hat)
        assert abs(inner(1, u_hat)-inner(1, ub_hat)) < tol
        T = TensorProductSpace(comm, (B0, B1))
        u_hat = Function(T)
        u_hat[1:4, 1:4] = 1
        Tp = T.get_refined((2*N[0], 2*N[1]))
        ub_hat = Function(Tp)
        u_hat.assign(ub_hat)
        assert abs(inner(1, u_hat)-inner(1, ub_hat)) < tol
        VT = VectorTensorProductSpace(T)
        u_hat = Function(VT)
        u_hat[:, 1:4, 1:4] = 1
        Tp = T.get_refined((2*N[0], 2*N[1]))
        VTp = VectorTensorProductSpace(Tp)
        ub_hat = Function(VTp)
        u_hat.assign(ub_hat)
        assert abs(inner((1, 1), u_hat)-inner((1, 1), ub_hat)) < tol
Example #4
0
def test_refine():
    assert comm.Get_size() < 7
    N = (8, 9, 10)
    F0 = Basis(8, 'F', dtype='D')
    F1 = Basis(9, 'F', dtype='D')
    F2 = Basis(10, 'F', dtype='d')
    T = TensorProductSpace(comm, (F0, F1, F2), slab=True, collapse_fourier=True)
    u_hat = Function(T)
    u = Array(T)
    u[:] = np.random.random(u.shape)
    u_hat = u.forward(u_hat)
    Tp = T.get_dealiased(padding_factor=(2, 2, 2))
    u_ = Array(Tp)
    up_hat = Function(Tp)
    assert up_hat.commsizes == u_hat.commsizes
    u2 = u_hat.refine(2*np.array(N))
    V = VectorTensorProductSpace(T)
    u_hat = Function(V)
    u = Array(V)
    u[:] = np.random.random(u.shape)
    u_hat = u.forward(u_hat)
    Vp = V.get_dealiased(padding_factor=(2, 2, 2))
    u_ = Array(Vp)
    up_hat = Function(Vp)
    assert up_hat.commsizes == u_hat.commsizes
    u3 = u_hat.refine(2*np.array(N))
Example #5
0
def test_curl_cc():
    theta, phi = sp.symbols('x,y', real=True, positive=True)
    psi = (theta, phi)
    r = 1
    rv = (r * sp.sin(theta) * sp.cos(phi), r * sp.sin(theta) * sp.sin(phi),
          r * sp.cos(theta))

    # Manufactured solution
    sph = sp.functions.special.spherical_harmonics.Ynm
    ue = sph(6, 3, theta, phi)

    N, M = 16, 12
    L0 = FunctionSpace(N, 'C', domain=(0, np.pi))
    F1 = FunctionSpace(M, 'F', dtype='D')
    T = TensorProductSpace(comm, (L0, F1), coordinates=(psi, rv))
    u_hat = Function(T, buffer=ue)
    du = curl(grad(u_hat))
    du.terms() == [[]]

    r, theta, z = psi = sp.symbols('x,y,z', real=True, positive=True)
    rv = (r * sp.cos(theta), r * sp.sin(theta), z)

    # Manufactured solution
    ue = (r * (1 - r) * sp.cos(4 * theta) - 1 * (r - 1)) * sp.cos(4 * z)

    N = 12
    F0 = FunctionSpace(N, 'F', dtype='D')
    F1 = FunctionSpace(N, 'F', dtype='d')
    L = FunctionSpace(N, 'L', bc='Dirichlet', domain=(0, 1))
    T = TensorProductSpace(comm, (L, F0, F1), coordinates=(psi, rv))
    T1 = T.get_orthogonal()
    V = VectorSpace(T1)
    u_hat = Function(T, buffer=ue)
    du = project(curl(grad(u_hat)), V)
    assert np.linalg.norm(du) < 1e-10
Example #6
0
def test_curl(typecode):
    K0 = Basis(N[0], 'F', dtype=typecode.upper())
    K1 = Basis(N[1], 'F', dtype=typecode.upper())
    K2 = Basis(N[2], 'F', dtype=typecode)
    T = TensorProductSpace(comm, (K0, K1, K2), dtype=typecode)
    X = T.local_mesh(True)
    K = T.local_wavenumbers()
    Tk = VectorTensorProductSpace(T)
    u = TrialFunction(Tk)
    v = TestFunction(Tk)

    U = Array(Tk)
    U_hat = Function(Tk)
    curl_hat = Function(Tk)
    curl_ = Array(Tk)

    # Initialize a Taylor Green vortex
    U[0] = np.sin(X[0]) * np.cos(X[1]) * np.cos(X[2])
    U[1] = -np.cos(X[0]) * np.sin(X[1]) * np.cos(X[2])
    U[2] = 0
    U_hat = Tk.forward(U, U_hat)
    Uc = U_hat.copy()
    U = Tk.backward(U_hat, U)
    U_hat = Tk.forward(U, U_hat)
    assert allclose(U_hat, Uc)

    divu_hat = project(div(U_hat), T)
    divu = Array(T)
    divu = T.backward(divu_hat, divu)
    assert allclose(divu, 0)

    curl_hat[0] = 1j * (K[1] * U_hat[2] - K[2] * U_hat[1])
    curl_hat[1] = 1j * (K[2] * U_hat[0] - K[0] * U_hat[2])
    curl_hat[2] = 1j * (K[0] * U_hat[1] - K[1] * U_hat[0])

    curl_ = Tk.backward(curl_hat, curl_)

    w_hat = Function(Tk)
    w_hat = inner(v, curl(U_hat), output_array=w_hat)
    A = inner(v, u)
    for i in range(3):
        w_hat[i] = A[i].solve(w_hat[i])

    w = Array(Tk)
    w = Tk.backward(w_hat, w)
    #from IPython import embed; embed()
    assert allclose(w, curl_)

    u_hat = Function(Tk)
    u_hat = inner(v, U, output_array=u_hat)
    for i in range(3):
        u_hat[i] = A[i].solve(u_hat[i])

    uu = Array(Tk)
    uu = Tk.backward(u_hat, uu)

    assert allclose(u_hat, U_hat)
Example #7
0
def test_backward2ND():
    T0 = FunctionSpace(N, 'C')
    L0 = FunctionSpace(N, 'L')
    T1 = FunctionSpace(N, 'C')
    L1 = FunctionSpace(N, 'L')
    TT = TensorProductSpace(comm, (T0, T1))
    LL = TensorProductSpace(comm, (L0, L1))
    uT = Function(TT, buffer=h)
    uL = Function(LL, buffer=h)
    u2 = uL.backward(kind=TT)
    uT2 = project(u2, TT)
    assert np.linalg.norm(uT2 - uT)
Example #8
0
def test_backward3D():
    T = FunctionSpace(N, 'C')
    L = FunctionSpace(N, 'L')
    F0 = FunctionSpace(N, 'F', dtype='D')
    F1 = FunctionSpace(N, 'F', dtype='d')
    TT = TensorProductSpace(comm, (F0, T, F1))
    TL = TensorProductSpace(comm, (F0, L, F1))
    uT = Function(TT, buffer=h)
    uL = Function(TL, buffer=h)

    u2 = uL.backward(kind=TT)
    uT2 = project(u2, TT)
    assert np.linalg.norm(uT2 - uT)
Example #9
0
def test_project(typecode, dim, ST, quad):
    # Using sympy to compute an analytical solution
    x, y, z = symbols("x,y,z")
    sizes = (25, 24)

    funcs = {
        (1, 0): (cos(4 * y) * sin(2 * np.pi * x)) * (1 - x**2),
        (1, 1): (cos(4 * x) * sin(2 * np.pi * y)) * (1 - y**2),
        (2, 0): (sin(6 * z) * cos(4 * y) * sin(2 * np.pi * x)) * (1 - x**2),
        (2, 1): (sin(2 * z) * cos(4 * x) * sin(2 * np.pi * y)) * (1 - y**2),
        (2, 2): (sin(2 * x) * cos(4 * y) * sin(2 * np.pi * z)) * (1 - z**2)
    }
    syms = {1: (x, y), 2: (x, y, z)}
    xs = {0: x, 1: y, 2: z}

    for shape in product(*([sizes] * dim)):
        bases = []
        for n in shape[:-1]:
            bases.append(Basis(n, 'F', dtype=typecode.upper()))
        bases.append(Basis(shape[-1], 'F', dtype=typecode))

        if dim < 3:
            n = min(shape)
            if typecode in 'fdg':
                n //= 2
                n += 1
            if n < comm.size:
                continue
        for axis in range(dim + 1):
            ST0 = ST(shape[-1], quad=quad)
            bases.insert(axis, ST0)
            fft = TensorProductSpace(comm,
                                     bases,
                                     dtype=typecode,
                                     axes=axes[dim][axis])
            X = fft.local_mesh(True)
            ue = funcs[(dim, axis)]
            ul = lambdify(syms[dim], ue, 'numpy')
            uq = ul(*X).astype(typecode)
            uh = Function(fft)
            uh = fft.forward(uq, uh)
            due = ue.diff(xs[axis], 1)
            dul = lambdify(syms[dim], due, 'numpy')
            duq = dul(*X).astype(typecode)
            uf = project(Dx(uh, axis, 1), fft)
            uy = Array(fft)
            uy = fft.backward(uf, uy)
            assert np.allclose(uy, duq, 0, 1e-6)
            for ax in (x for x in range(dim + 1) if x is not axis):
                due = ue.diff(xs[axis], 1, xs[ax], 1)
                dul = lambdify(syms[dim], due, 'numpy')
                duq = dul(*X).astype(typecode)
                uf = project(Dx(Dx(uh, axis, 1), ax, 1), fft)
                uy = Array(fft)
                uy = fft.backward(uf, uy)
                assert np.allclose(uy, duq, 0, 1e-6)

            bases.pop(axis)
            fft.destroy()
Example #10
0
 def __init__(self, comm, N):
     method_common.__init__(self, comm, N)
     # initial spectral spaces
     K0 = C2CBasis(N[0], domain=(-2 * np.pi, 2 * np.pi))
     K1 = R2CBasis(N[1], domain=(-2 * np.pi, 2 * np.pi))
     self.T = TensorProductSpace(comm, (K0, K1),
                                 **{'planner_effort': 'FFTW_MEASURE'})
Example #11
0
def test_helmholtz2D(family, axis):
    la = lla
    if family == 'chebyshev':
        la = cla
    N = (8, 9)
    SD = FunctionSpace(N[axis], family=family, bc=(0, 0))
    K1 = FunctionSpace(N[(axis+1)%2], family='F', dtype='d')
    subcomms = mpi4py_fft.pencil.Subcomm(MPI.COMM_WORLD, allaxes2D[axis])
    bases = [K1]
    bases.insert(axis, SD)
    T = TensorProductSpace(subcomms, bases, axes=allaxes2D[axis], modify_spaces_inplace=True)
    u = shenfun.TrialFunction(T)
    v = shenfun.TestFunction(T)
    if family == 'chebyshev':
        mat = inner(v, div(grad(u)))
    else:
        mat = inner(grad(v), grad(u))

    H = la.Helmholtz(*mat)
    u = Function(T)
    s = SD.sl[SD.slice()]
    u[s] = np.random.random(u[s].shape) + 1j*np.random.random(u[s].shape)
    f = Function(T)
    f = H.matvec(u, f)

    g0 = Function(T)
    g1 = Function(T)
    M = {d.get_key(): d for d in mat}
    g0 = M['ADDmat'].matvec(u, g0)
    g1 = M['BDDmat'].matvec(u, g1)
    assert np.linalg.norm(f-(g0+g1)) < 1e-12, np.linalg.norm(f-(g0+g1))

    uc = Function(T)
    uc = H(uc, f)
    assert np.linalg.norm(uc-u) < 1e-12
Example #12
0
def test_mixed_2D(backend, forward_output, as_scalar):
    if (backend == 'netcdf4' and forward_output is True) or skip[backend]:
        return
    K0 = FunctionSpace(N[0], 'F')
    K1 = FunctionSpace(N[1], 'C')
    T = TensorProductSpace(comm, (K0, K1))
    TT = CompositeSpace([T, T])
    filename = 'test2Dm_{}'.format(ex[forward_output])
    hfile = writer(filename, TT, backend=backend)
    if forward_output:
        uf = Function(TT, val=2)
    else:
        uf = Array(TT, val=2)
    hfile.write(0, {'uf': [uf]}, as_scalar=as_scalar)
    hfile.write(1, {'uf': [uf]}, as_scalar=as_scalar)
    if not forward_output and backend == 'hdf5' and comm.Get_rank() == 0:
        generate_xdmf(filename + '.h5')
    if as_scalar is False:
        u0 = Function(TT) if forward_output else Array(TT)
        read = reader(filename, TT, backend=backend)
        read.read(u0, 'uf', step=1)
        assert np.allclose(u0, uf)
    else:
        u0 = Function(T) if forward_output else Array(T)
        read = reader(filename, T, backend=backend)
        read.read(u0, 'uf0', step=1)
        assert np.allclose(u0, uf[0])
    cleanup()
Example #13
0
def test_biharmonic2D(family, axis):
    la = lla
    if family == 'chebyshev':
        la = cla
    N = (16, 16)
    SD = FunctionSpace(N[axis], family=family, bc='Biharmonic')
    K1 = FunctionSpace(N[(axis + 1) % 2], family='F', dtype='d')
    subcomms = mpi4py_fft.pencil.Subcomm(MPI.COMM_WORLD, allaxes2D[axis])
    bases = [K1]
    bases.insert(axis, SD)
    T = TensorProductSpace(subcomms, bases, axes=allaxes2D[axis])
    u = shenfun.TrialFunction(T)
    v = shenfun.TestFunction(T)
    if family == 'chebyshev':
        mat = inner(v, div(grad(div(grad(u)))))
    else:
        mat = inner(div(grad(v)), div(grad(u)))

    H = la.Biharmonic(*mat)
    u = Function(T)
    u[:] = np.random.random(u.shape) + 1j * np.random.random(u.shape)
    f = Function(T)
    f = H.matvec(u, f)

    g0 = Function(T)
    g1 = Function(T)
    g2 = Function(T)
    M = {d.get_key(): d for d in mat}
    g0 = M['SBBmat'].matvec(u, g0)
    g1 = M['ABBmat'].matvec(u, g1)
    g2 = M['BBBmat'].matvec(u, g2)

    assert np.linalg.norm(f - (g0 + g1 + g2)) < 1e-8
Example #14
0
def test_regular_3D(backend, forward_output):
    if (backend == 'netcdf4' and forward_output is True) or skip[backend]:
        return
    K0 = FunctionSpace(N[0], 'F', dtype='D', domain=(0, np.pi))
    K1 = FunctionSpace(N[1], 'F', dtype='d', domain=(0, 2 * np.pi))
    K2 = FunctionSpace(N[2], 'C', dtype='d', bc=(0, 0))
    T = TensorProductSpace(comm, (K0, K1, K2))
    filename = 'test3Dr_{}'.format(ex[forward_output])
    hfile = writer(filename, T, backend=backend)
    if forward_output:
        u = Function(T)
    else:
        u = Array(T)
    u[:] = np.random.random(u.shape)
    for i in range(3):
        hfile.write(i, {
            'u':
            [u, (u, [slice(None), 4, slice(None)]), (u, [slice(None), 4, 4])]
        },
                    forward_output=forward_output)
    if not forward_output and backend == 'hdf5' and comm.Get_rank() == 0:
        generate_xdmf(filename + '.h5')

    u0 = Function(T) if forward_output else Array(T)
    read = reader(filename, T, backend=backend)
    read.read(u0, 'u', forward_output=forward_output, step=1)
    assert np.allclose(u0, u)
    cleanup()
Example #15
0
def test_eval_expression():
    import sympy as sp
    from shenfun import div, grad
    x, y, z = sp.symbols('x,y,z')
    B0 = Basis(16, 'C')
    B1 = Basis(17, 'C')
    B2 = Basis(20, 'F', dtype='d')

    TB = TensorProductSpace(comm, (B0, B1, B2))

    f = sp.sin(x)+sp.sin(y)+sp.sin(z)
    dfx = f.diff(x, 2) + f.diff(y, 2) + f.diff(z, 2)
    fa = Array(TB, buffer=f).forward()
    dfe = div(grad(fa))
    dfa = project(dfe, TB)

    xyz = np.array([[0.25, 0.5, 0.75],
                    [0.25, 0.5, 0.75],
                    [0.25, 0.5, 0.75]])

    f0 = lambdify((x, y, z), dfx)(*xyz)
    f1 = dfe.eval(xyz)
    f2 = dfa.eval(xyz)
    assert np.allclose(f0, f1, 1e-7)
    assert np.allclose(f1, f2, 1e-7)
Example #16
0
def test_mixed_3D(backend, forward_output, as_scalar):
    if (backend == 'netcdf4' and forward_output is True) or skip[backend]:
        return
    K0 = FunctionSpace(N[0], 'F', dtype='D', domain=(0, np.pi))
    K1 = FunctionSpace(N[1], 'F', dtype='d', domain=(0, 2 * np.pi))
    K2 = FunctionSpace(N[2], 'C')
    T = TensorProductSpace(comm, (K0, K1, K2))
    TT = VectorSpace(T)
    filename = 'test3Dm_{}'.format(ex[forward_output])
    hfile = writer(filename, TT, backend=backend)
    uf = Function(TT, val=2) if forward_output else Array(TT, val=2)
    uf[0] = 1
    data = {
        'ux': (uf[0], (uf[0], [slice(None), 4,
                               slice(None)]), (uf[0], [slice(None), 4, 4])),
        'uy': (uf[1], (uf[1], [slice(None), 4,
                               slice(None)]), (uf[1], [slice(None), 4, 4])),
        'u': [uf, (uf, [slice(None), 4, slice(None)])]
    }
    hfile.write(0, data, as_scalar=as_scalar)
    hfile.write(1, data, as_scalar=as_scalar)
    if not forward_output and backend == 'hdf5' and comm.Get_rank() == 0:
        generate_xdmf(filename + '.h5')

    if as_scalar is False:
        u0 = Function(TT) if forward_output else Array(TT)
        read = reader(filename, TT, backend=backend)
        read.read(u0, 'u', step=1)
        assert np.allclose(u0, uf)
    else:
        u0 = Function(T) if forward_output else Array(T)
        read = reader(filename, T, backend=backend)
        read.read(u0, 'u0', step=1)
        assert np.allclose(u0, uf[0])
    cleanup()
Example #17
0
def test_curl2():
    # Test projection of curl

    K0 = FunctionSpace(N[0], 'C', bc=(0, 0))
    K1 = FunctionSpace(N[1], 'F', dtype='D')
    K2 = FunctionSpace(N[2], 'F', dtype='d')
    K3 = FunctionSpace(N[0], 'C')

    T = TensorProductSpace(comm, (K0, K1, K2))
    TT = TensorProductSpace(comm, (K3, K1, K2))
    X = T.local_mesh(True)
    K = T.local_wavenumbers(False)
    Tk = VectorSpace(T)
    TTk = VectorSpace([T, T, TT])

    U = Array(Tk)
    U_hat = Function(Tk)
    curl_hat = Function(TTk)
    curl_ = Array(TTk)

    # Initialize a Taylor Green vortex
    U[0] = np.sin(X[0]) * np.cos(X[1]) * np.cos(X[2]) * (1 - X[0]**2)
    U[1] = -np.cos(X[0]) * np.sin(X[1]) * np.cos(X[2]) * (1 - X[0]**2)
    U[2] = 0
    U_hat = Tk.forward(U, U_hat)
    Uc = U_hat.copy()
    U = Tk.backward(U_hat, U)
    U_hat = Tk.forward(U, U_hat)
    assert allclose(U_hat, Uc)

    # Compute curl first by computing each term individually
    curl_hat[0] = 1j * (K[1] * U_hat[2] - K[2] * U_hat[1])
    curl_[0] = T.backward(
        curl_hat[0], curl_[0])  # No x-derivatives, still in Dirichlet space
    dwdx_hat = project(Dx(U_hat[2], 0, 1), TT)  # Need to use space without bc
    dvdx_hat = project(Dx(U_hat[1], 0, 1), TT)  # Need to use space without bc
    dwdx = Array(TT)
    dvdx = Array(TT)
    dwdx = TT.backward(dwdx_hat, dwdx)
    dvdx = TT.backward(dvdx_hat, dvdx)
    curl_hat[1] = 1j * K[2] * U_hat[0]
    curl_hat[2] = -1j * K[1] * U_hat[0]
    curl_[1] = T.backward(curl_hat[1], curl_[1])
    curl_[2] = T.backward(curl_hat[2], curl_[2])
    curl_[1] -= dwdx
    curl_[2] += dvdx

    # Now do it with project
    w_hat = project(curl(U_hat), TTk)
    w = Array(TTk)
    w = TTk.backward(w_hat, w)
    assert allclose(w, curl_)
Example #18
0
def test_project2(typecode, dim, ST, quad):
    # Using sympy to compute an analytical solution
    x, y, z = symbols("x,y,z")
    sizes = (18, 17)

    funcx = ((2*np.pi**2*(x**2 - 1) - 1)* cos(2*np.pi*x) - 2*np.pi*x*sin(2*np.pi*x))/(4*np.pi**3)
    funcy = ((2*np.pi**2*(y**2 - 1) - 1)* cos(2*np.pi*y) - 2*np.pi*y*sin(2*np.pi*y))/(4*np.pi**3)
    funcz = ((2*np.pi**2*(z**2 - 1) - 1)* cos(2*np.pi*z) - 2*np.pi*z*sin(2*np.pi*z))/(4*np.pi**3)

    funcs = {
        (1, 0): cos(4*y)*funcx,
        (1, 1): cos(4*x)*funcy,
        (2, 0): sin(3*z)*cos(4*y)*funcx,
        (2, 1): sin(2*z)*cos(4*x)*funcy,
        (2, 2): sin(2*x)*cos(4*y)*funcz
        }
    syms = {1: (x, y), 2:(x, y, z)}
    xs = {0:x, 1:y, 2:z}

    for shape in product(*([sizes]*dim)):
        bases = []
        for n in shape[:-1]:
            bases.append(Basis(n, 'F', dtype=typecode.upper()))
        bases.append(Basis(shape[-1], 'F', dtype=typecode))

        for axis in range(dim+1):
            ST0 = ST(shape[-1], quad=quad)
            bases.insert(axis, ST0)
            # Spectral space must be aligned in nonperiodic direction, hence axes
            fft = TensorProductSpace(comm, bases, dtype=typecode, axes=axes[dim][axis])
            X = fft.local_mesh(True)
            ue = funcs[(dim, axis)]
            ul = lambdify(syms[dim], ue, 'numpy')
            uq = ul(*X).astype(typecode)
            uh = Function(fft)
            uh = fft.forward(uq, uh)
            due = ue.diff(xs[axis], 1)
            dul = lambdify(syms[dim], due, 'numpy')
            duq = dul(*X).astype(typecode)
            uf = project(Dx(uh, axis, 1), fft)
            uy = Array(fft)
            uy = fft.backward(uf, uy)
            assert np.allclose(uy, duq, 0, 1e-3)

            # Test also several derivatives
            for ax in (x for x in range(dim+1) if x is not axis):
                due = ue.diff(xs[ax], 1, xs[axis], 1)
                dul = lambdify(syms[dim], due, 'numpy')
                duq = dul(*X).astype(typecode)
                uf = project(Dx(Dx(uh, ax, 1), axis, 1), fft)
                uy = Array(fft)
                uy = fft.backward(uf, uy)
                assert np.allclose(uy, duq, 0, 1e-3)
            bases.pop(axis)
            fft.destroy()
Example #19
0
def test_shentransform(typecode, dim, ST, quad):

    for shape in product(*([sizes] * dim)):
        bases = []
        for n in shape[:-1]:
            bases.append(Basis(n, 'F', dtype=typecode.upper()))
        bases.append(Basis(shape[-1], 'F', dtype=typecode))

        if dim < 3:
            n = min(shape)
            if typecode in 'fdg':
                n //= 2
                n += 1
            if n < comm.size:
                continue
        for axis in range(dim + 1):
            ST0 = ST(shape[-1], quad=quad)
            bases.insert(axis, ST0)
            fft = TensorProductSpace(comm, bases, dtype=typecode)
            U = random_like(fft.forward.input_array)
            F = fft.forward(U)
            Fc = F.copy()
            V = fft.backward(F)
            F = fft.forward(U)
            assert allclose(F, Fc)
            bases.pop(axis)
            fft.destroy()
def main(N, family, bci, bcj, plotting=False):
    global fe, ue
    BX = FunctionSpace(N, family=family, bc=bcx[bci], domain=xdomain)
    BY = FunctionSpace(N, family=family, bc=bcy[bcj], domain=ydomain)
    T = TensorProductSpace(comm, (BX, BY))
    u = TrialFunction(T)
    v = TestFunction(T)

    # Get f on quad points
    fj = Array(T, buffer=fe)

    # Compare with analytical solution
    ua = Array(T, buffer=ue)

    if T.use_fixed_gauge:
        mean = dx(ua, weighted=True) / inner(1, Array(T, val=1))

    # Compute right hand side of Poisson equation
    f_hat = Function(T)
    f_hat = inner(v, fj, output_array=f_hat)

    # Get left hand side of Poisson equation
    A = inner(v, -div(grad(u)))

    u_hat = Function(T)

    sol = la.Solver2D(A, fixed_gauge=mean if T.use_fixed_gauge else None)
    u_hat = sol(f_hat, u_hat)
    uj = u_hat.backward()

    assert np.allclose(uj, ua), np.linalg.norm(uj - ua)
    print("Error=%2.16e" % (np.sqrt(dx((uj - ua)**2))))

    if 'pytest' not in os.environ and plotting is True:
        import matplotlib.pyplot as plt
        X, Y = T.local_mesh(True)
        plt.contourf(X, Y, uj, 100)
        plt.colorbar()
        plt.figure()
        plt.contourf(X, Y, ua, 100)
        plt.colorbar()
        plt.figure()
        plt.contourf(X, Y, ua - uj, 100)
        plt.colorbar()
Example #21
0
def get_function_space(space='cylinder'):
    if space == 'cylinder':
        r, theta, z = psi = sp.symbols('x,y,z', real=True, positive=True)
        rv = (r*sp.cos(theta), r*sp.sin(theta), z)
        N = 6
        F0 = FunctionSpace(N, 'F', dtype='D')
        F1 = FunctionSpace(N, 'F', dtype='d')
        L = FunctionSpace(N, 'L', domain=(0, 1))
        T = TensorProductSpace(comm, (L, F0, F1), coordinates=(psi, rv))

    elif space == 'sphere':
        r, theta, phi = psi = sp.symbols('x,y,z', real=True, positive=True)
        rv = (r*sp.sin(theta)*sp.cos(phi), r*sp.sin(theta)*sp.sin(phi), r*sp.cos(theta))
        N = 6
        F = FunctionSpace(N, 'F', dtype='d')
        L0 = FunctionSpace(N, 'L', domain=(0, 1))
        L1 = FunctionSpace(N, 'L', domain=(0, np.pi))
        T = TensorProductSpace(comm, (L0, L1, F), coordinates=(psi, rv, sp.Q.positive(sp.sin(theta))))

    return T
Example #22
0
def test_backward2D():
    T = FunctionSpace(N, 'C')
    L = FunctionSpace(N, 'L')
    F = FunctionSpace(N, 'F', dtype='d')
    TT = TensorProductSpace(comm, (T, F))
    TL = TensorProductSpace(comm, (L, F))
    uT = Function(TT, buffer=f)
    uL = Function(TL, buffer=f)

    u2 = uL.backward(kind=TT)
    uT2 = project(u2, TT)
    assert np.linalg.norm(uT2 - uT)

    TT = TensorProductSpace(comm, (F, T))
    TL = TensorProductSpace(comm, (F, L))
    uT = Function(TT, buffer=f)
    uL = Function(TL, buffer=f)

    u2 = uL.backward(kind=TT)
    uT2 = project(u2, TT)
    assert np.linalg.norm(uT2 - uT)
Example #23
0
def test_padding_neumann(family):
    N = 8
    B = FunctionSpace(N, family, bc={'left': ('N', 0), 'right': ('N', 0)})
    Bp = B.get_dealiased(1.5)
    u = Function(B)
    u[1:-2] = np.random.random(N - 3)
    up = Array(Bp)
    up = Bp.backward(u, fast_transform=False)
    uf = Bp.forward(up, fast_transform=False)
    assert np.linalg.norm(uf - u) < 1e-12
    if family == 'C':
        up = Bp.backward(u)
        uf = Bp.forward(up)
        assert np.linalg.norm(uf - u) < 1e-12

    # Test padding 2D
    F = FunctionSpace(N, 'F', dtype='d')
    T = TensorProductSpace(comm, (B, F))
    Tp = T.get_dealiased(1.5)
    u = Function(T)
    u[1:-2, :-1] = np.random.random(u[1:-2, :-1].shape)
    up = Tp.backward(u)
    uc = Tp.forward(up)
    assert np.linalg.norm(u - uc) < 1e-8

    # Test padding 3D
    F1 = FunctionSpace(N, 'F', dtype='D')
    T = TensorProductSpace(comm, (F1, F, B))
    Tp = T.get_dealiased(1.5)
    u = Function(T)
    u[:, :, 1:-2] = np.random.random(u[:, :, 1:-2].shape)
    u = u.backward().forward()  # Clean
    up = Tp.backward(u)
    uc = Tp.forward(up)
    assert np.linalg.norm(u - uc) < 1e-8
Example #24
0
def test_padding_biharmonic(family):
    N = 8
    B = FunctionSpace(N, family, bc=(0, 0, 0, 0))
    Bp = B.get_dealiased(1.5)
    u = Function(B)
    u[:(N - 4)] = np.random.random(N - 4)
    up = Array(Bp)
    up = Bp.backward(u, fast_transform=False)
    uf = Bp.forward(up, fast_transform=False)
    assert np.linalg.norm(uf - u) < 1e-12
    if family == 'C':
        up = Bp.backward(u)
        uf = Bp.forward(up)
        assert np.linalg.norm(uf - u) < 1e-12

    # Test padding 2D
    F = FunctionSpace(N, 'F', dtype='d')
    T = TensorProductSpace(comm, (B, F))
    Tp = T.get_dealiased(1.5)
    u = Function(T)
    u[:-4, :-1] = np.random.random(u[:-4, :-1].shape)
    up = Tp.backward(u)
    uc = Tp.forward(up)
    assert np.linalg.norm(u - uc) < 1e-8

    # Test padding 3D
    F1 = FunctionSpace(N, 'F', dtype='D')
    T = TensorProductSpace(comm, (F1, F, B))
    Tp = T.get_dealiased(1.5)
    u = Function(T)
    u[:, :, :-4] = np.random.random(u[:, :, :-4].shape)
    u = u.backward().forward()  # Clean
    up = Tp.backward(u)
    uc = Tp.forward(up)
    assert np.linalg.norm(u - uc) < 1e-8
def test_project_2dirichlet(quad):
    x, y = symbols("x,y")
    ue = (cos(4*y)*sin(2*x))*(1-x**2)*(1-y**2)
    sizes = (18, 17)

    D0 = lbases.ShenDirichlet(sizes[0], quad=quad)
    D1 = lbases.ShenDirichlet(sizes[1], quad=quad)
    B0 = lbases.Orthogonal(sizes[0], quad=quad)
    B1 = lbases.Orthogonal(sizes[1], quad=quad)

    DD = TensorProductSpace(comm, (D0, D1))
    BD = TensorProductSpace(comm, (B0, D1))
    DB = TensorProductSpace(comm, (D0, B1))
    BB = TensorProductSpace(comm, (B0, B1))

    X = DD.local_mesh(True)
    uh = Function(DD, buffer=ue)
    dudx_hat = project(Dx(uh, 0, 1), BD)
    dx = Function(BD, buffer=ue.diff(x, 1))
    assert np.allclose(dx, dudx_hat, 0, 1e-5)

    dudy = project(Dx(uh, 1, 1), DB).backward()
    duedy = Array(DB, buffer=ue.diff(y, 1))
    assert np.allclose(duedy, dudy, 0, 1e-5)

    us_hat = Function(BB)
    uq = uh.backward()
    us = project(uq, BB, output_array=us_hat).backward()
    assert np.allclose(us, uq, 0, 1e-5)
    dudxy = project(Dx(us_hat, 0, 1) + Dx(us_hat, 1, 1), BB).backward()
    dxy = Array(BB, buffer=ue.diff(x, 1) + ue.diff(y, 1))
    assert np.allclose(dxy, dudxy, 0, 1e-5), np.linalg.norm(dxy-dudxy)
Example #26
0
def test_inner(f0, f1):
    if f0 == 'F' and f1 == 'F':
        B0 = Basis(8, f0, dtype='D', domain=(-2*np.pi, 2*np.pi))
    else:
        B0 = Basis(8, f0)
    c = Array(B0, val=1)
    d = inner(1, c)
    assert abs(d-(B0.domain[1]-B0.domain[0])) < 1e-7

    B1 = Basis(8, f1)
    T = TensorProductSpace(comm, (B0, B1))
    a0 = Array(T, val=1)
    c0 = inner(1, a0)
    L = np.array([b.domain[1]-b.domain[0] for b in (B0, B1)])
    assert abs(c0-np.prod(L)) < 1e-7

    if not (f0 == 'F' or f1 == 'F'):
        B2 = Basis(8, f1, domain=(-2, 2))
        T = TensorProductSpace(comm, (B0, B1, B2))
        a0 = Array(T, val=1)
        c0 = inner(1, a0)
        L = np.array([b.domain[1]-b.domain[0] for b in (B0, B1, B2)])
        assert abs(c0-np.prod(L)) < 1e-7
Example #27
0
def test_Mult_CTD_3D(quad):
    SD = FunctionSpace(N, 'C', bc=(0, 0))
    F0 = FunctionSpace(4, 'F', dtype='D')
    F1 = FunctionSpace(4, 'F', dtype='d')
    T = TensorProductSpace(comm, (SD, F0, F1))

    TO = T.get_orthogonal()
    CT = TO.bases[0]

    C = inner_product((CT, 0), (SD, 1))
    B = inner_product((CT, 0), (CT, 0))

    vk = Array(T)
    wk = Array(T)
    vk[:] = np.random.random(vk.shape)
    wk[:] = np.random.random(vk.shape)

    bv = Function(T)
    bw = Function(T)

    vk0 = vk.forward()
    vk = vk0.backward()
    wk0 = wk.forward()
    wk = wk0.backward()

    LUsolve.Mult_CTD_3D_ptr(N, vk0, wk0, bv, bw, 0)

    cv = np.zeros_like(vk0)
    cw = np.zeros_like(wk0)
    cv = C.matvec(vk0, cv)
    cw = C.matvec(wk0, cw)
    cv /= B[0].repeat(np.array(bv.shape[1:]).prod()).reshape(bv.shape)
    cw /= B[0].repeat(np.array(bv.shape[1:]).prod()).reshape(bv.shape)

    assert np.allclose(cv, bv)
    assert np.allclose(cw, bw)
Example #28
0
def test_energy_fourier(N):
    B0 = FunctionSpace(N[0], 'F', dtype='D')
    B1 = FunctionSpace(N[1], 'F', dtype='D')
    B2 = FunctionSpace(N[2], 'F', dtype='d')
    for bases, axes in zip(((B0, B1, B2), (B0, B2, B1)),
                           ((0, 1, 2), (2, 0, 1))):
        T = TensorProductSpace(comm, bases, axes=axes)
        u_hat = Function(T)
        u_hat[:] = np.random.random(
            u_hat.shape) + 1j * np.random.random(u_hat.shape)
        u = u_hat.backward()
        u_hat = u.forward(u_hat)
        u = u_hat.backward(u)
        e0 = comm.allreduce(np.sum(u.v * u.v) / np.prod(N))
        e1 = fourier.energy_fourier(u_hat, T)
        assert abs(e0 - e1) < 1e-10
Example #29
0
def test_eval_fourier(typecode, dim):
    # Using sympy to compute an analytical solution
    x, y, z = symbols("x,y,z")
    sizes = (20, 19)

    funcs = {
        2: cos(4 * x) + sin(6 * y),
        3: sin(6 * x) + cos(4 * y) + sin(8 * z)
    }
    syms = {2: (x, y), 3: (x, y, z)}
    points = None
    if comm.Get_rank() == 0:
        points = np.random.random((dim, 3))
    points = comm.bcast(points)
    t_0 = 0
    t_1 = 0
    t_2 = 0
    for shape in product(*([sizes] * dim)):
        bases = []
        for n in shape[:-1]:
            bases.append(Basis(n, 'F', dtype=typecode.upper()))

        bases.append(Basis(shape[-1], 'F', dtype=typecode))
        fft = TensorProductSpace(comm, bases, dtype=typecode)
        X = fft.local_mesh(True)
        ue = funcs[dim]
        ul = lambdify(syms[dim], ue, 'numpy')
        uu = ul(*X).astype(typecode)
        uq = ul(*points).astype(typecode)
        u_hat = Function(fft)
        u_hat = fft.forward(uu, u_hat)
        t0 = time()
        result = fft.eval(points, u_hat, method=0)
        t_0 += time() - t0
        assert allclose(uq, result), print(uq, result)
        t0 = time()
        result = fft.eval(points, u_hat, method=1)
        t_1 += time() - t0
        assert allclose(uq, result)
        t0 = time()
        result = fft.eval(points, u_hat, method=2)
        t_2 += time() - t0
        assert allclose(uq, result), print(uq, result)

    print('method=0', t_0)
    print('method=1', t_1)
    print('method=2', t_2)
Example #30
0
def test_regular_2D(backend, forward_output):
    if (backend == 'netcdf4' and forward_output is True) or skip[backend]:
        return
    K0 = FunctionSpace(N[0], 'F')
    K1 = FunctionSpace(N[1], 'C', bc=(0, 0))
    T = TensorProductSpace(comm, (K0, K1))
    filename = 'test2Dr_{}'.format(ex[forward_output])
    hfile = writer(filename, T, backend=backend)
    u = Function(T, val=1) if forward_output else Array(T, val=1)
    hfile.write(0, {'u': [u]}, forward_output=forward_output)
    hfile.write(1, {'u': [u]}, forward_output=forward_output)
    if not forward_output and backend == 'hdf5' and comm.Get_rank() == 0:
        generate_xdmf(filename + '.h5')
        generate_xdmf(filename + '.h5', order='visit')

    u0 = Function(T) if forward_output else Array(T)
    read = reader(filename, T, backend=backend)
    read.read(u0, 'u', forward_output=forward_output, step=1)
    assert np.allclose(u0, u)
Example #31
0
def get_context():
    """Set up context for classical (NS) solver"""
    float, complex, mpitype = datatypes(params.precision)
    collapse_fourier = False if params.dealias == '3/2-rule' else True
    dim = len(params.N)
    dtype = lambda d: float if d == dim-1 else complex
    V = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
               dtype=dtype(i)) for i in range(dim)]

    kw0 = {'threads': params.threads,
           'planner_effort': params.planner_effort['fft']}
    T = TensorProductSpace(comm, V, dtype=float,
                           slab=(params.decomposition == 'slab'),
                           collapse_fourier=collapse_fourier, **kw0)
    VT = VectorTensorProductSpace(T)

    # Different bases for nonlinear term, either 2/3-rule or 3/2-rule
    kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
          'dealias_direct': params.dealias == '2/3-rule'}

    Vp = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
                dtype=dtype(i), **kw) for i in range(dim)]

    Tp = TensorProductSpace(comm, Vp, dtype=float,
                            slab=(params.decomposition == 'slab'),
                            collapse_fourier=collapse_fourier, **kw0)
    VTp = VectorTensorProductSpace(Tp)

    # Mesh variables
    X = T.local_mesh(True)
    K = T.local_wavenumbers(scaled=True)
    for i in range(dim):
        X[i] = X[i].astype(float)
        K[i] = K[i].astype(float)
    K2 = np.zeros(T.shape(True), dtype=float)
    for i in range(dim):
        K2 += K[i]*K[i]

    # Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
    Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
    for i in range(dim):
        Kx[i] = Kx[i].astype(float)

    K_over_K2 = np.zeros(VT.shape(True), dtype=float)
    for i in range(dim):
        K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)

    # Velocity and pressure. Use ndarray view for efficiency
    U = Array(VT)
    U_hat = Function(VT)
    P = Array(T)
    P_hat = Function(T)
    u_dealias = Array(VTp)

    # Primary variable
    u = U_hat

    # RHS array
    dU = Function(VT)
    curl = Array(VT)
    Source = Function(VT) # Possible source term initialized to zero
    work = work_arrays()

    hdf5file = NSFile(config.params.solver,
                      checkpoint={'space': VT,
                                  'data': {'0': {'U': [U_hat]}}},
                      results={'space': VT,
                               'data': {'U': [U], 'P': [P]}})

    return config.AttributeDict(locals())
Example #32
0
def get_context():
    float, complex, mpitype = datatypes(params.precision)
    collapse_fourier = False if params.dealias == '3/2-rule' else True
    dim = len(params.N)
    dtype = lambda d: float if d == dim-1 else complex
    V = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
               dtype=dtype(i)) for i in range(dim)]

    kw0 = {'threads': params.threads,
           'planner_effort': params.planner_effort['fft']}
    T = TensorProductSpace(comm, V, dtype=float,
                           slab=(params.decomposition == 'slab'),
                           collapse_fourier=collapse_fourier, **kw0)
    VT = VectorTensorProductSpace(T)
    VM = MixedTensorProductSpace([T]*2*dim)

    kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
          'dealias_direct': params.dealias == '2/3-rule'}

    Vp = [Basis(params.N[i], 'F', domain=(0, params.L[i]),
                dtype=dtype(i), **kw) for i in range(dim)]

    Tp = TensorProductSpace(comm, Vp, dtype=float,
                            slab=(params.decomposition == 'slab'),
                            collapse_fourier=collapse_fourier, **kw0)
    VTp = VectorTensorProductSpace(Tp)
    VMp = MixedTensorProductSpace([Tp]*2*dim)

    # Mesh variables
    X = T.local_mesh(True)
    K = T.local_wavenumbers(scaled=True)
    for i in range(dim):
        X[i] = X[i].astype(float)
        K[i] = K[i].astype(float)
    K2 = np.zeros(T.shape(True), dtype=float)
    for i in range(dim):
        K2 += K[i]*K[i]

    # Set Nyquist frequency to zero on K that is, from now on, used for odd derivatives
    Kx = T.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
    for i in range(dim):
        Kx[i] = Kx[i].astype(float)

    K_over_K2 = np.zeros(VT.shape(True), dtype=float)
    for i in range(dim):
        K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)

    UB = Array(VM)
    P = Array(T)
    curl = Array(VT)
    UB_hat = Function(VM)
    P_hat = Function(T)
    dU = Function(VM)
    Source = Array(VM)
    ub_dealias = Array(VMp)
    ZZ_hat = np.zeros((3, 3) + Tp.shape(True), dtype=complex) # Work array

    # Create views into large data structures
    U = UB[:3]
    U_hat = UB_hat[:3]
    B = UB[3:]
    B_hat = UB_hat[3:]

    # Primary variable
    u = UB_hat

    hdf5file = MHDFile(config.params.solver,
                       checkpoint={'space': VM,
                                   'data': {'0': {'UB': [UB_hat]}}},
                       results={'space': VM,
                                'data': {'UB': [UB]}})

    return config.AttributeDict(locals())
Example #33
0
def get_context():
    """Set up context for solver"""

    # Get points and weights for Chebyshev weighted integrals
    assert params.Dquad == params.Bquad
    collapse_fourier = False if params.dealias == '3/2-rule' else True
    ST = Basis(params.N[2], 'C', bc=(0, 0), quad=params.Dquad)
    SB = Basis(params.N[2], 'C', bc='Biharmonic', quad=params.Bquad)
    CT = Basis(params.N[2], 'C', quad=params.Dquad)
    ST0 = Basis(params.N[2], 'C', bc=(0, 0), quad=params.Dquad) # For 1D problem
    K0 = Basis(params.N[0], 'F', domain=(0, params.L[0]), dtype='D')
    K1 = Basis(params.N[1], 'F', domain=(0, params.L[1]), dtype='d')

    kw0 = {'threads':params.threads,
           'planner_effort':params.planner_effort["dct"],
           'slab': (params.decomposition == 'slab'),
           'collapse_fourier': collapse_fourier}
    FST = TensorProductSpace(comm, (K0, K1, ST), axes=(2, 0, 1), **kw0)    # Dirichlet
    FSB = TensorProductSpace(comm, (K0, K1, SB), axes=(2, 0, 1), **kw0)    # Biharmonic
    FCT = TensorProductSpace(comm, (K0, K1, CT), axes=(2, 0, 1), **kw0)    # Regular Chebyshev
    VFS = MixedTensorProductSpace([FST, FST, FSB])
    VFST = MixedTensorProductSpace([FST, FST, FST])
    VUG = MixedTensorProductSpace([FST, FSB])

    # Padded
    kw = {'padding_factor': 1.5 if params.dealias == '3/2-rule' else 1,
          'dealias_direct': params.dealias == '2/3-rule'}
    if params.dealias == '3/2-rule':
        # Requires new bases due to planning and transforms on different size arrays
        STp = Basis(params.N[2], 'C', bc=(0, 0), quad=params.Dquad)
        SBp = Basis(params.N[2], 'C', bc='Biharmonic', quad=params.Bquad)
        CTp = Basis(params.N[2], 'C', quad=params.Dquad)
    else:
        STp, SBp, CTp = ST, SB, CT
    K0p = Basis(params.N[0], 'F', dtype='D', domain=(0, params.L[0]), **kw)
    K1p = Basis(params.N[1], 'F', dtype='d', domain=(0, params.L[1]), **kw)
    FSTp = TensorProductSpace(comm, (K0p, K1p, STp), axes=(2, 0, 1), **kw0)
    FSBp = TensorProductSpace(comm, (K0p, K1p, SBp), axes=(2, 0, 1), **kw0)
    FCTp = TensorProductSpace(comm, (K0p, K1p, CTp), axes=(2, 0, 1), **kw0)
    VFSp = MixedTensorProductSpace([FSTp, FSTp, FSBp])

    Nu = params.N[2]-2   # Number of velocity modes in Shen basis
    Nb = params.N[2]-4   # Number of velocity modes in Shen biharmonic basis
    u_slice = slice(0, Nu)
    v_slice = slice(0, Nb)

    float, complex, mpitype = datatypes("double")

    # Mesh variables
    X = FST.local_mesh(True)
    x0, x1, x2 = FST.mesh()
    K = FST.local_wavenumbers(scaled=True)

    # Solution variables
    U = Array(VFS)
    U0 = Array(VFS)
    U_hat = Function(VFS)
    U_hat0 = Function(VFS)
    g = Function(FST)

    # primary variable
    u = (U_hat, g)

    H_hat = Function(VFST)
    H_hat0 = Function(VFST)
    H_hat1 = Function(VFST)

    dU = Function(VUG)
    hv = Function(FST)
    hg = Function(FST)
    Source = Array(VFS)
    Sk = Function(VFS)

    K2 = K[0]*K[0]+K[1]*K[1]
    K4 = K2**2

    # Set Nyquist frequency to zero on K that is used for odd derivatives in nonlinear terms
    Kx = FST.local_wavenumbers(scaled=True, eliminate_highest_freq=True)
    K_over_K2 = np.zeros((2,)+g.shape)
    for i in range(2):
        K_over_K2[i] = K[i] / np.where(K2 == 0, 1, K2)

    work = work_arrays()
    u_dealias = Array(VFSp)
    u0_hat = np.zeros((2, params.N[2]), dtype=complex)
    h0_hat = np.zeros((2, params.N[2]), dtype=complex)
    w = np.zeros((params.N[2], ), dtype=complex)
    w1 = np.zeros((params.N[2], ), dtype=complex)

    nu, dt, N = params.nu, params.dt, params.N

    # Collect all matrices
    mat = config.AttributeDict(
        dict(CDD=inner_product((ST, 0), (ST, 1)),
             CTD=inner_product((CT, 0), (ST, 1)),
             BTT=inner_product((CT, 0), (CT, 0)),
             AB=HelmholtzCoeff(N[2], 1.0, -(K2 - 2.0/nu/dt), 2, ST.quad),
             AC=BiharmonicCoeff(N[2], nu*dt/2., (1. - nu*dt*K2), -(K2 - nu*dt/2.*K4), 2, SB.quad),
             # Matrices for biharmonic equation
             CBD=inner_product((SB, 0), (ST, 1)),
             ABB=inner_product((SB, 0), (SB, 2)),
             BBB=inner_product((SB, 0), (SB, 0)),
             SBB=inner_product((SB, 0), (SB, 4)),
             # Matrices for Helmholtz equation
             ADD=inner_product((ST, 0), (ST, 2)),
             BDD=inner_product((ST, 0), (ST, 0)),
             BBD=inner_product((SB, 0), (ST, 0)),
             CDB=inner_product((ST, 0), (SB, 1)),
             ADD0=inner_product((ST0, 0), (ST0, 2)),
             BDD0=inner_product((ST0, 0), (ST0, 0))))

    la = config.AttributeDict(
        dict(HelmholtzSolverG=Helmholtz(mat.ADD, mat.BDD, -np.ones((1, 1, 1)),
                                        (K2+2.0/nu/dt)),
             BiharmonicSolverU=Biharmonic(mat.SBB, mat.ABB, mat.BBB, -nu*dt/2.*np.ones((1, 1, 1)),
                                          (1.+nu*dt*K2),
                                          (-(K2 + nu*dt/2.*K4))),
             HelmholtzSolverU0=Helmholtz(mat.ADD0, mat.BDD0, np.array([-1.]), np.array([2./nu/dt])),
             TDMASolverD=TDMA(inner_product((ST, 0), (ST, 0)))))

    hdf5file = KMMFile(config.params.solver,
                       checkpoint={'space': VFS,
                                   'data': {'0': {'U': [U_hat]},
                                            '1': {'U': [U_hat0]}}},
                       results={'space': VFS,
                                'data': {'U': [U]}})

    return config.AttributeDict(locals())