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 bcs
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 initial viscosity guess
nu_guess = 0.2

# Define actual viscosity
def nu(u):
    return 0.5*pow(inner(grad(u), grad(u)), 1.0/(2*(4-1)))

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

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

# Plot initial guesses
print "||u_guess||_0^2 = ", assemble(inner(u, u)*dx)
plot(u, title="Velocity initial guess")
plot(p, title="Pressure initial guess")

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

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

# Plot solutions
print "||u||_0^2 = ", assemble(inner(u, u)*dx)
plot(u, title="Velocity")
plot(p, title="Pressure")
interactive()