from mpi4py import MPI
from petsc4py import PETSc

rank = MPI.COMM_WORLD.rank

def test(dof=1, width=1):
  da = PETSc.DA().create(dim=2, sizes=[8, 8], dof=dof,
                         boundary_type='periodic',
                         stencil_width=width, 
                         stencil_type='box')

  gvec = da.createGlobalVec()
  lvec = da.createLocalVector()

  with da.getVecArray(lvec) as va:
    if dof == 1:
      va.array[width:-width,width:-width] = rank+1
    else:
      va.array[width:-width,width:-width,0] = rank+1
  da.localToGlobal(lvec, gvec)

  da.globalToLocal(gvec, lvec)
  with da.getVecArray(lvec) as va:
    if dof == 1:
      print va.array[:,:]
    else:
      print va.array[:,:,0]

if rank==0: print 'DOF=1:' 
test(dof=1, width=1)

MPI.COMM_WORLD.Barrier()

if rank==0: print 'DOF=3:'
test(dof=3, width=1)