*scratch*.py

from firedrake import *
from firedrake.petsc import PETSc
# Mixed poisson
# Homogeneous dirichlet BCs on all walls
# Forcing
# f = -\pi^2 (4 cos(\pi x) - 5 cos(\pi x/2) + 2) sin(\pi y)
# With exact solution
# sin(\pi x) tan(\pi x/4) sin(\pi y)

mesh = UnitSquareMesh(100, 100)

V = FunctionSpace(mesh, "RT", 1)
Q = FunctionSpace(mesh, "DG", 0)

W = V*Q

sigma, u = TrialFunctions(W)
tau, v = TestFunctions(W)

f = Function(Q)

f.interpolate(Expression("-0.5*pi*pi*(4*cos(pi*x[0]) - 5*cos(pi*x[0]*0.5) + 2)*sin(pi*x[1])"))
a = (dot(sigma, tau) + div(tau)*u + div(sigma)*v)*dx

n = FacetNormal(mesh)

L = -f*v*dx

exact = Function(Q)

exact.interpolate(Expression("sin(pi*x[0])*tan(pi*x[0]*0.25)*sin(pi*x[1])"))


solution = Function(W)

alpha = Constant(4.0)
gamma = Constant(8.0)
h = CellSize(mesh)
h_avg = (h('+') + h('-'))/2
a_dg = dot(sigma, tau)*dx  + \
       dot(grad(v), grad(u))*dx \
       - dot(avg(grad(v)), jump(u, n))*dS \
       - dot(jump(v, n), avg(grad(u)))*dS \
       + alpha/h_avg * dot(jump(v, n), jump(u, n))*dS \
       - dot(grad(v), u*n)*ds \
       - dot(v*n, grad(u))*ds \
       + (gamma/h)*dot(v, u)*ds

# aP is an optional form that is assembled and then used by PETSc to
# build the preconditioning matrices
problem_ip = LinearVariationalProblem(a, L, solution,
                                      aP=a_dg)

ip_parameters = {
    'ksp_type': 'gmres',
    'ksp_monitor_true_residual': True,
    # Upper Schur factorisation
    # Precondition S = D - C A^{-1} B with an interior penalty DG
    # Laplacian (since it's spectrally equivalent)
    'pc_type': 'fieldsplit',
    'pc_fieldsplit_type': 'schur',
    'pc_fieldsplit_schur_fact_type': 'upper',
    # a11 is the IP-DG block from aP
    'pc_fieldsplit_schur_precondition': 'a11',
    # Approximately invert A with a single application of
    # process-block ILU(0)
    'fieldsplit_0_ksp_type': 'preonly',
    'fieldsplit_0_pc_type': 'bjacobi',
    'fieldsplit_0_sub_pc_type': 'ilu',
    # Approximately invert S with a single V-cycle on Sp
    # This means we never apply the action of the schur complement and
    # let the outer GMRES iteration fix things up
    'fieldsplit_1_ksp_type': 'preonly',
    'fieldsplit_1_pc_type': 'hypre'
    }
    
solver_ip = LinearVariationalSolver(problem_ip,
                                    options_prefix="ip_",
                                    solver_parameters=ip_parameters)

with PETSc.Log.Stage("dg_ip"):
    solver_ip.solve()

sigma, u = solution.split()

print norm(assemble(u - exact))