from firedrake import *



# Create mesh and define function space

mesh = UnitSquareMesh(64, 64)

V = FunctionSpace(mesh, "CG", 1)

R = FunctionSpace(mesh, "R", 0)

W = V*R



# Define variational problem

u, c = TrialFunctions(W)

v, d = TestFunctions(W)

f = interpolate(Expression("10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)"),W.sub(0))

g = interpolate(Expression("-sin(5*x[0])"),W.sub(0))



a = (inner(grad(u), grad(v)) + c*v + u*d)*dx

L = f*v*dx + g*v*ds



# Compute solution

w = Function(W)

solve(a == L, w)



u, c = w.split()



outfile = File("solution.pvd")

outfile.write(u)



z = Function(R)



print z.function_space()

print c.function_space()



z.assign(c)