Compound bases

dedalus/tests/test_lbvp.py

@@ -98,6 +98,49 @@ def test_poisson_1d_jacobi(Nx, a0, b0, da, db, dtype):
     assert np.allclose(u['g'], u_true)
 
 
+@pytest.mark.parametrize('Nx', [32])
+@pytest.mark.parametrize('a0', [-1/2, 0])
+@pytest.mark.parametrize('b0', [-1/2, 0])
+@pytest.mark.parametrize('da', [0, 1])
+@pytest.mark.parametrize('db', [0, 1])
+@pytest.mark.parametrize('dtype', dtype_range)
+def test_poisson_1d_compound(Nx, a0, b0, da, db, dtype):
+    # Bases
+    coord = d3.Coordinate('x')
+    dist = d3.Distributor(coord, dtype=dtype)
+    b1 = d3.Jacobi(coord, size=Nx, bounds=(0, np.pi), a=a0+da, b=b0+db, a0=a0, b0=b0)
+    b2 = d3.Jacobi(coord, size=Nx, bounds=(np.pi, 2*np.pi), a=a0+da, b=b0+db, a0=a0, b0=b0)
+    basis = d3.Compound((b1, b2))
+    x = basis.local_grid(1)
+    # Fields
+    u = dist.Field(name='u', bases=basis)
+    tau11 = dist.Field(name='tau11')
+    tau12 = dist.Field(name='tau12')
+    tau21 = dist.Field(name='tau21')
+    tau22 = dist.Field(name='tau22')
+    tau1 = (tau11, tau12)
+    tau2 = (tau21, tau22)
+    # Substitutions
+    dx = lambda A: d3.Differentiate(A, coord)
+    lift_basis = basis.derivative_basis(2)
+    lift = lambda A, n: d3.Lift(A, lift_basis, n)
+    F = dist.Field(bases=basis)
+    F['g'] = -np.sin(x)
+    # Problem
+    problem = d3.LBVP([u, *tau1, *tau2], namespace=locals())
+    problem.add_equation("dx(dx(u)) + lift(tau1,-1) + lift(tau2,-2) = F")
+    problem.add_equation("jump(u, x=np.pi) = 0")
+    problem.add_equation("jump(dx(u), x=np.pi) = 0")
+    problem.add_equation("u(x='left') = 0")
+    problem.add_equation("u(x='right') = 0")
+    # Solver
+    solver = problem.build_solver()
+    solver.solve()
+    # Check solution
+    u_true = np.sin(x)
+    assert np.allclose(u['g'], u_true)
+
+
 radius_disk = 1