Hi, you need to modify the criteria that was used to identify the boundary cells in the post you referred to. Here are two examples of such logic for the unit square an x=1 as the boundary
from dolfin import *
mesh = UnitSquareMesh(10, 10)
# a cell with at least one vertex on boundary belongs to boundary
boundary_cells = [cell for cell in cells(mesh)\
if any([near(vertex.x(0), 1) for vertex in vertices(cell)])]
# a cell with at least one facet with midpoint on boundary belongs to boundary
#boundary_cells = [cell for cell in cells(mesh)\
# if any([near(facet.midpoint().x(), 1) for facet in facets(cell)])]
all_cells = CellFunction("size_t", mesh, 0)
for cell in boundary_cells:
all_cells[cell] = 1
plot(all_cells, interactive=True)