from dolfin import *

mesh = UnitSquareMesh(64, 64)
V = FunctionSpace(mesh, 'CG', 1)

# define subdomain 1
cf = CellFunction('size_t', mesh, 0)
subd1 = AutoSubDomain(lambda x : x[1] >= x[0] - DOLFIN_EPS)
subd1.mark(cf, 1)

# apply dofmap
dofmap = V.dofmap()
subd0_dofs = []
subd1_dofs = []

# get dofs of subdomains
for cell in cells(mesh): 
    if cf[cell] == 0:
        subd0_dofs.extend(dofmap.cell_dofs(cell.index()))
    else:
        subd1_dofs.extend(dofmap.cell_dofs(cell.index()))

subd0_dofs = list(set(subd0_dofs))
subd1_dofs = list(set(subd1_dofs))

# set functions for the 2 subdomains
u = interpolate(Expression("0"), V)
v = interpolate(Expression("1"), V)
u1 = Function(V)

u1.vector()[subd0_dofs] = u.vector()[subd0_dofs]
u1.vector()[subd1_dofs] = v.vector()[subd1_dofs]

plot(u1, interactive=True)