# Copyright (C) 2007 Kristian B. Oelgaard
#
# This demo solves the time-dependent convection-diffusion equation 

from dolfin import *

# Sub domain for Dirichlet boundary condition
class DirichletBoundary(SubDomain):
    def inside(self, x, on_boundary):
        return bool(on_boundary and
                    x[0] < 1.0 - DOLFIN_EPS and
                    x[0] > 0.0 + DOLFIN_EPS and
                    x[1] > 0.0 + DOLFIN_EPS and
                    x[1] < 1.0 - DOLFIN_EPS)

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

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

# Create velocity Function from file
velocity = Function(V);
File("velocity.xml.gz") >> velocity

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

# Parameters
T = 5.0
dt = 0.1
t = dt
c = 0.00005

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

# Variational problem for cG(1)dG(0) for convection-diffusion
a = v*u*dx + dt*(v*dot(velocity, grad(u))*dx + c*dot(grad(v), grad(u))*dx)
L = v*u0*dx - dt*v*f*dx

# Set up boundary condition
boundary = DirichletBoundary()
g = Constant( 1.0 )
zero = Constant( 0.0 )
bc = DirichletBC(Q, g, sub_domains, 1)
bc2 = DirichletBC(Q, zero, boundary)

# Assemble matrix
A = assemble(a)
bc.apply(A)

# Create linear solver and factorize matrix
solver = LUSolver(A)
solver.parameters["reuse_factorization"] = True

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

# Set intial condition
u = u0

# Time-stepping
while t < T:

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

    # Solve the linear system (re-use the already factorized matrix A)
    solver.solve(u.vector(), b)

    # Copy solution from previous interval
    u0 = u

    # Plot solution
    plot(u)

    # Save the solution to file
    out_file << (u, t)

    # Move to next interval
    t += dt

# Hold plot
#interactive()