from firedrake import *

nx = 10
ny = 10
nz = 5

mesh = BoxMesh(nx, ny, nz, 2.0, 2.0, 1.0)

V = FunctionSpace(mesh, 'Lagrange', 1)

p = Function(V, name="p")
phi = Function(V, name="phi")

u = TrialFunction(V)
v = TestFunction(V)

outfile = File("out.pvd")
outfile << phi

bcval = Constant(0.0)
bc = DirichletBC(V, bcval, 1)

T = 10.
dt = 0.001
t = 0
step = 0
lump_mass = True

while t <= T:
    step += 1
    bcval.assign(sin(2*pi*5*t))
    phi -= dt / 2 * p
    if lump_mass:
        p += assemble(dt * inner(nabla_grad(v), nabla_grad(phi))*dx) / assemble(v*dx)
        bc.apply(p)
    else:
        solve(u*v*dx == v*p*dx + dt*inner(grad(v), grad(phi))*dx, p, bcs=bc, solver_parameters={'ksp_type': 'cg', 'pc_type': 'sor', 'pc_sor_symmetric': True})

	phi -= dt / 2 * p
    t += dt
    if step % 10 == 0:
        outfile << phi