from dolfin import *

mesh = Mesh("cyl.xml")

V = VectorFunctionSpace(mesh, 'CG', 1)
u = TrialFunction(V)
v = TestFunction(V)

# Lag SubDomains av grenseflatene i topp og bunn
top = AutoSubDomain(lambda x, on_bnd: near(x[2], 10.))
bottom = AutoSubDomain(lambda x, on_bnd: near(x[2], 0.))

# Marker grenseflatene med en FacetFunction 
mf = FacetFunction('uint', mesh)
mf.set_all(0)
top.mark(mf, 1)
bottom.mark(mf, 2)

# Lag en grensebetingelse for P*n
traction = Constant((0, 0, -1))
#traction = Expression(("0", "0", "std::abs(x[0]) < 1e-1 && std::abs(x[1]) < 1e-1 ? -1 : 0"))

# Lag en grensebetingelse for ingen forskyvning i topp av cylinderen
bc1 = DirichletBC(V, Constant((0, 0, 0)), top)

# Sett konstanter
Lambda = Constant(20)
mu = Constant(10)

# Sett opp likningene 
P = Lambda * div(u) * Identity(3) + mu * (grad(u) + grad(u).T)
F = inner(grad(v), P)*dx - inner(v, traction)*ds(2)

# Assemble and solve problem
A = assemble(lhs(F), exterior_facet_domains=mf)
b = assemble(rhs(F), exterior_facet_domains=mf)
u = Function(V)
bc1.apply(A, b)
solve(A, u.vector(), b)
plot(u[2], title="Forskyvning i z-retning", interactive=True)