from dolfin import *

# Use -02 optimization
parameters["form_compiler"]["cpp_optimize"] = True

# Define mesh and geometry
mesh = Mesh("dolfin-2.xml.gz")
n = FacetNormal(mesh)

# Define Taylor--Hood function space W
V = VectorFunctionSpace(mesh, "CG" , 2)
Q = FunctionSpace(mesh , "CG", 1)
W = V * Q

# Define Function and TestFunction(s)
w = Function(W)
(u, p) = split(w)
(v, q) = TestFunctions(W)

# Define viscosity and bcs
nu = Expression("0.2*(1+pow(x[1],2))", degree=2)
p0 = Expression("1.0-x[0]", degree=1)
bcs = DirichletBC(W.sub(0), (0.0, 0.0),
                  "on_boundary && !(near(x[0], 0.0) || near(x[0], 1.0))")

# Define variational form
epsilon = sym(grad(u))
F = (nu*inner(epsilon, grad(v)) - div(u)*q - div(v)*p)*dx\
    + p0*dot(v,n)*ds

# Solve problem
solve(F == 0, w, bcs)

# Plot solutions
plot(u, title="Velocity", interactive=True)
plot(p, title="Pressure", interactive=True)