diff --git a/neumann_bc.py b/neumann_bc.py
index c78223c..04fab2f 100644
--- a/neumann_bc.py
+++ b/neumann_bc.py
@@ -10,38 +10,52 @@
 from ufl import grad, inner, div, dot
 from mpi4py import MPI
 from petsc4py import PETSc
-from utils import norm_L2
-from dolfinx.fem.petsc import assemble_matrix, assemble_vector
+from utils import norm_L2, markers_to_meshtags
+from dolfinx.fem.petsc import assemble_matrix, assemble_vector, apply_lifting
 
 
-def boundary_marker(x):
-    "A function to mark the domain boundary"
+def u_e_expr(x, module=np):
+    "Expression for the exact solution"
+    return module.sin(module.pi * x[0]) * module.sin(module.pi * x[1])
 
-    return (
-        np.isclose(x[0], 0.0)
-        | np.isclose(x[0], 1.0)
-        | np.isclose(x[1], 0.0)
-        | np.isclose(x[1], 1.0)
-    )
 
+def f_expr(x):
+    "Source term"
+    return 2 * np.pi**2 * u_e_expr(x)
 
-# Create a mesh and a sub-mesh of the boundary
+
+def g_expr(x):
+    "Neumann boundary condition"
+    return np.pi * np.cos(np.pi * x[0]) * np.sin(np.pi * x[1])
+
+
+# Create a mesh and meshtags for the Dirichlet and Neumann boundaries
 n = 8
 msh = mesh.create_unit_square(MPI.COMM_WORLD, n, n)
+
 tdim = msh.topology.dim
 fdim = tdim - 1
-num_facets = msh.topology.create_entities(fdim)
-boundary_facets = mesh.locate_entities_boundary(msh, fdim, boundary_marker)
-submesh, submesh_to_mesh = mesh.create_submesh(msh, fdim, boundary_facets)[0:2]
+boundaries = {"dirichlet": 1, "neumann": 2}  # Tags for the boundaries
+markers = [
+    lambda x: np.isclose(x[0], 0.0) | np.isclose(x[1], 0.0) | np.isclose(x[1], 1.0),
+    lambda x: np.isclose(x[0], 1.0),
+]
+ft = markers_to_meshtags(msh, boundaries.values(), markers, fdim)
+
+# Create a submesh of the Neumann boundary
+neumann_boundary_facets = ft.find(boundaries["neumann"])
+submesh, submesh_to_mesh = mesh.create_submesh(msh, fdim, neumann_boundary_facets)[:2]
 
 # Create function spaces
 k = 3  # Polynomial degree
 V = fem.functionspace(msh, ("Lagrange", k))
-W = fem.functionspace(submesh, ("Lagrange", k))
+# Function space for the Neumann boundary condition. Note that this is defined
+# only over the Neumann boundary
+W = fem.functionspace(submesh, ("Lagrange", k))  
 u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
 
 # Create integration measure and entity maps
-ds = ufl.Measure("ds", domain=msh)
+ds = ufl.Measure("ds", domain=msh, subdomain_data=ft)
 # We take msh to be the integration domain, so we must provide a map from
 # facets in msh to cells in submesh. This is simply the "inverse" of
 # submesh_to_mesh
@@ -51,24 +65,34 @@ def boundary_marker(x):
 msh_to_submesh[submesh_to_mesh] = np.arange(len(submesh_to_mesh))
 entity_maps = {submesh: msh_to_submesh}
 
-# Create manufactured solution
-x = ufl.SpatialCoordinate(msh)
-u_e = 1
-for i in range(tdim):
-    u_e *= ufl.sin(ufl.pi * x[i])
-f = u_e - div(grad(u_e))
-n = ufl.FacetNormal(msh)
-g = dot(grad(u_e), n)
+# Interpolate the source term and Neumann boundary condition
+f = fem.Function(V)
+f.interpolate(f_expr)
+
+g = fem.Function(W)
+g.interpolate(g_expr)
 
-# Define forms
-a = fem.form(inner(u, v) * ufl.dx + inner(grad(u), grad(v)) * ufl.dx)
-L = fem.form(inner(f, v) * ufl.dx + inner(g, v) * ds, entity_maps=entity_maps)
+# Define forms. Since the Neumann boundary term involves funcriotns defined over
+# different meshes, we must provide entity maps
+a = fem.form(inner(grad(u), grad(v)) * ufl.dx)
+L = fem.form(
+    inner(f, v) * ufl.dx + inner(g, v) * ds(boundaries["neumann"]), entity_maps=entity_maps
+)
+
+# Dirichlet boundary condition
+dirichlet_facets = ft.find(boundaries["dirichlet"])
+dirichlet_dofs = fem.locate_dofs_topological(V, fdim, dirichlet_facets)
+u_d = fem.Function(V)
+u_d.interpolate(u_e_expr)
+bc = fem.dirichletbc(u_d, dirichlet_dofs)
 
 # Assemble matrix and vector
-A = assemble_matrix(a)
+A = assemble_matrix(a, bcs=[bc])
 A.assemble()
 b = assemble_vector(L)
+apply_lifting(b, [a], bcs=[[bc]])
 b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
+bc.set(b.array_w)
 
 # Create solver
 ksp = PETSc.KSP().create(msh.comm)
@@ -87,7 +111,8 @@ def boundary_marker(x):
     f.write(0.0)
 
 # Compute the error
-e_L2 = norm_L2(msh.comm, u - u_e)
+x = ufl.SpatialCoordinate(msh)
+e_L2 = norm_L2(msh.comm, u - u_e_expr(x, module=ufl))
 
 if msh.comm.rank == 0:
     print(f"e_L2 = {e_L2}")
diff --git a/utils.py b/utils.py
index 75062ac..8a7a363 100644
--- a/utils.py
+++ b/utils.py
@@ -343,6 +343,16 @@ def one_sided_int_entities(msh, facets):
 
 
 def markers_to_meshtags(msh, tags, markers, dim):
+    """
+    Create meshtags from a given set of tags and markers. Tag[i] is is the tag for
+    the part of the domain marked by markers[i].\
+
+    Parameters:
+        msh: The mesh
+        tags: A list of tags
+        markers: A list of markers
+        dim: The dimension of the entities
+    """
     entities = [mesh.locate_entities_boundary(msh, dim, marker) for marker in markers]
     values = [np.full_like(entities, tag) for (tag, entities) in zip(tags, entities)]
     entities = np.hstack(entities, dtype=np.int32)