*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)

# If boundary conditions are such that div-div is a norm, rather than
# a semi-norm don't need the mass term in the H(div) part
a_riesz = u*v*dx + div(sigma)*div(tau)*dx + dot(sigma, tau)*dx

problem_riesz = LinearVariationalProblem(a, L, solution,
                                         aP=a_riesz)

riesz_parameters = {
    'ksp_type': 'gmres',
    'ksp_monitor_true_residual': True,
    # Additive fieldsplit
    # Precondition by the H(div)-L2 inner product
    'pc_type': 'fieldsplit',
    'pc_fieldsplit_type': 'additive',
    # Invert H(div) part with direct solver
    # Could hook in to HYPRE's auxiliary space ADS multigrid (PETSc
    # support is being developed by Stefano Zampini, but we'd need to
    # add some hooks)
    'fieldsplit_0_ksp_type': 'preonly',
    'fieldsplit_0_pc_type': 'lu',
    'fieldsplit_0_pc_factor_mat_solver_package': 'mumps',
    # 1-1 block is just a DG mass matrix, so invert exactly with
    # process-block ILU(0)
    'fieldsplit_1_ksp_type': 'preonly',
    'fieldsplit_1_pc_type': 'bjacobi',
    'fieldsplit_1_sub_pc_type': 'ilu',
    }
solver_riesz = LinearVariationalSolver(problem_riesz,
                                       options_prefix="riesz_",
                                       solver_parameters=riesz_parameters)


with PETSc.Log.Stage("riesz"):
    solver_riesz.solve()

sigma, u = solution.split()

print norm(assemble(u - exact))