from firedrake import *
from firedrake.petsc import PETSc
import firedrake.cython.dmcommon as dmplex
import numpy as np


mesh = UnitCubeMesh(20, 20, 20)
PETSc.Sys.Print('\nDM after Create:')
mesh.topology_dm.view()

if True: # without the following code, everything works fine in parallel.
    x = SpatialCoordinate(mesh)
    V = TensorFunctionSpace(mesh, 'CG', 1)
    c = 10
    f = Function(V, name='f').interpolate(as_matrix([[c, 0, 0], [0, c, 0], [0, 0, c]]))
    PETSc.Sys.Print('\nDM after Init:')
    mesh.topology_dm.view()
    File('old2.pvd').write(f)


plex = mesh.topology_dm
new_plex = plex.coarsen()

new_mesh = Mesh(new_plex)
x = SpatialCoordinate(new_mesh)
V = FunctionSpace(new_mesh, 'CG', 1)
f = Function(V, name='f').interpolate(10 + 10*sin(x[0]))
File('new_mesh.pvd').write(f)