from dolfin import *

# Load mesh and subdomains
mesh = Mesh("mesh.xml.gz")
sub_domains = MeshFunction("uint", mesh, "subdomains.xml.gz");

# Create FunctionSpaces
Q = FunctionSpace(mesh, "CG", 1)

# Initialise source function and previous solution function
f  = Function(Q, "0.0")
u0 = Function(Q)

# Parameters
T = 5.0
k = 0.1
t = k
c = 0.02

# Test and trial functions
v = TestFunction(Q)
u = TrialFunction(Q)

# Functions
u1 = Function(Q)

# Variational problem
a = v*u*dx + 0.5*k*c*dot(grad(v), grad(u))*dx
L = v*u0*dx + k*v*f*dx

# Set up boundary condition
g = Constant(mesh, 1.0)
bc = DirichletBC(Q, g, sub_domains, 1)

# Assemble matrix
A = assemble(a)

# Output file
out_file = File("temperature.pvd")

# Time-stepping
while t < T:

    # Copy solution from previous interval
    u0.assign(u1)

    # Assemble vector and apply boundary conditions
    b = assemble(L)
    bc.apply(A, b)

    # Solve the linear system
    solve(A, u1.vector(), b)

    # Plot solution
    plot(u1)

    # Save the solution to file
    out_file << u1

    # Move to next interval
    t += k

# Hold plot
interactive()