from dolfin import *
from IPython.display import Image

mesh = UnitSquareMesh(6, 6)

wiz = plot(mesh)
wiz.write_png("mesh")
Image("mesh.png")

V = FunctionSpace(mesh, 'Lagrange', 1)

u0 = Expression('1 + x[0]*x[0] + 2*x[1]*x[1]')

def u0_boundary(x, on_boundary):
    return on_boundary

bc = DirichletBC(V, u0, u0_boundary)

u = TrialFunction(V)
v = TestFunction(V)

f = Constant(50.0)
a = inner(nabla_grad(u), nabla_grad(v))*dx
L = f*v*dx

# Compute solution
u = Function(V)
solve(a == L, u, bc)

# Plot solution and mesh
wiz = plot(u)
wiz.write_png("u")
Image("u.png")

dx?