How to use the petsc4py.PETSc.COMM_WORLD function in petsc4py

# encoding: UTF-8
from __future__ import absolute_import, division, print_function, unicode_literals
from petsc4py import PETSc


def f(snes, x, f):
    for i in range(N):
        f[i] = (x[i] ** 2 - 9.0).item()
    f[N // 2] = x[N // 2]
    f.assemble()

COMM = PETSc.COMM_WORLD
N = 4
J = PETSc.Mat().createAIJ(N, comm=COMM)
J.setPreallocationNNZ(N)
for i in range(N):
    for j in range(N):
        J.setValue(i, j, 0.0)
J.setUp()
J.assemble()

dm = PETSc.DMShell().create(comm=COMM)
dm.setMatrix(J)

snes = PETSc.SNES().create(comm=COMM)
r = PETSc.Vec().createSeq(N)
x = PETSc.Vec().createSeq(N)
b = PETSc.Vec().createSeq(N)

Istart, Iend = A.getOwnershipRange()
for I in range(Istart, Iend):
    A.setValue(I, I, diagv)
    i = I // n  # map row number to
    j = I - i * n  # grid coordinates
    if i > 0: J = I - n; A.setValues(I, J, offdx)
    if i < n - 1: J = I + n; A.setValues(I, J, offdx)
    if j > 0: J = I - 1; A.setValues(I, J, offdy)
    if j < n - 1: J = I + 1; A.setValues(I, J, offdy)
A.assemblyBegin()
A.assemblyEnd()

A.mult(x, b)

rank = PETSc.COMM_WORLD.getRank()
print(rank, b.getArray())
print((b-y).norm(PETSc.NormType.NORM_INFINITY))

def main():
    # import petsc4py


    n = 4
    dx = 1.0/(n - 1)
    dy = dx
    comm= PETSc.COMM_WORLD
    da = PETSc.DMDA().create([n, n], dof=1, stencil_width=1, comm=comm)
    dar = da.refine()
    print(dar.getSizes())
    exit()

    rank = PETSc.COMM_WORLD.getRank()
    # comm=

    x = da.createGlobalVec()
    xa = da.getVecArray(x)
    (xs, xe), (ys, ye) = da.getRanges()
    print(xs,xe,ys,ye, xa.shape)
    for i in range(xs, xe):
        for j in range(ys, ye):
            xa[i, j, 0] = np.sin(2 * np.pi * (i ) * dx) * np.sin(2 * np.pi * (j ) * dy)
            xa[i, j, 1] = 0.1 * np.sin(2 * np.pi * (i ) * dx) * np.sin(2 * np.pi * (j ) * dy)

def _solve_KSP(self, guess, matrix, vector1, vector2):
        """
        Set PETSC KSP solver.

        Args
            guess: Boolean specifying if the iterative KSP solver initial guess is nonzero
            matrix: PETSC matrix used by the KSP solver
            vector1: PETSC vector corresponding to the initial values
            vector2: PETSC vector corresponding to the new values

        Returns:
            vector2: PETSC vector of the new values
        """

        # guess = False
        ksp = PETSc.KSP().create(PETSc.COMM_WORLD)
        if guess:
            ksp.setInitialGuessNonzero(guess)
        ksp.setOperators(matrix,matrix)
        ksp.setType('richardson')
        pc = ksp.getPC()
        pc.setType('bjacobi')
        ksp.setTolerances(rtol=self.rtol)
        ksp.solve(vector1, vector2)
        r = ksp.getConvergedReason()
        if r < 0:
            KSPReasons = self._make_reasons(PETSc.KSP.ConvergedReason())
            print('LinearSolver failed to converge after %d iterations',ksp.getIterationNumber())
            print('with reason: %s',KSPReasons[r])
            raise RuntimeError("LinearSolver failed to converge!")
        ksp.destroy()

pc.setType('none')
ksp.setFromOptions()

ba = da.getVecArray(b)
mx, my = da.getSizes()
hx, hy = [1.0/m for m in [mx, my]]
(xs, xe), (ys, ye) = da.getRanges()
for i in range(xs, xe):
    for j in range(ys, ye):
        ba[i, j] = 1*hx*hy
# pde.formRHS(b)
# b.setUp()
# print(PETSc.COMM_WORLD.getRank(), b.getArray())
ksp.solve(b, x)

print(PETSc.COMM_WORLD.getRank(), x.getArray())

u = da.createNaturalVec()
da.globalToNatural(x, u)

print('nat', PETSc.COMM_WORLD.getRank(), u.getArray())

def arrayToVec(vecArray):
    """ Converts a (global) array to a PETSc vector over :attr:`petsc4py.PETSc.COMM_WORLD`.

        Args:
            vecArray (array or numpy.array): input vector.

        Returns:
            petsc4py.PETSc.Vec() :
        """
    vec = _PETSc.Vec().create(comm=_PETSc.COMM_WORLD)
    vec.setSizes(len(vecArray))
    vec.setUp()
    (Istart,Iend) = vec.getOwnershipRange()
    return vec.createWithArray(vecArray[Istart:Iend],
            comm=_PETSc.COMM_WORLD)
    vec.destroy()

def call_in_rank_order(fun, comm=None):
    """
    Call a function `fun` task by task in the task rank order.
    """
    if comm is None:
        comm = PETSc.COMM_WORLD

    for rank in range(comm.size):
        if rank == comm.rank:
            fun(rank, comm)
        comm.barrier()

def __init__(self,ghosted_csr_mat=None,par_bs=None,par_n=None,par_N=None,par_nghost=None,subdomain2global=None,blockVecType="simple",pde=None, par_nc=None, par_Nc=None, proteus_jacobian=None, nzval_proteus2petsc=None):
        p4pyPETSc.Mat.__init__(self)
        if ghosted_csr_mat == None:
            return#when duplicating for petsc usage
        self.pde = pde
        if par_nc == None:
            par_nc = par_n
        if par_Nc == None:
            par_Nc = par_N
        self.proteus_jacobian=proteus_jacobian
        self.nzval_proteus2petsc = nzval_proteus2petsc
        self.ghosted_csr_mat=ghosted_csr_mat
        self.blockVecType = blockVecType
        assert self.blockVecType == "simple", "petsc4py wrappers require self.blockVecType=simple"
        self.create(p4pyPETSc.COMM_WORLD)
        self.blockSize = max(1,par_bs)
        if self.blockSize > 1 and blockVecType != "simple":
            ## \todo fix block aij in ParMat_petsc4py
            self.setType('baij')
            self.setSizes([[self.blockSize*par_n,self.blockSize*par_N],[self.blockSize*par_nc,self.blockSize*par_Nc]],bsize=self.blockSize)
            self.setBlockSize(self.blockSize)
            self.subdomain2global = subdomain2global #no need to include extra block dofs?
        else:
            self.setType('aij')
            self.setSizes([[par_n*self.blockSize,par_N*self.blockSize],[par_nc*self.blockSize,par_Nc*self.blockSize]],bsize=1)
            if self.blockSize > 1: #have to build in block dofs
                subdomain2globalTotal = numpy.zeros((self.blockSize*subdomain2global.shape[0],),'i')
                for j in range(self.blockSize):
                    subdomain2globalTotal[j::self.blockSize]=subdomain2global*self.blockSize+j
                self.subdomain2global=subdomain2globalTotal
            else:

def add_history(nc):
    "Update the history attribute in a NetCDF file nc."
    comm = PETSc.COMM_WORLD
    rank = comm.getRank()

    if rank != 0:
        return

    # add history global attribute (after checking if present)
    historysep = ' '
    historystr = asctime() + ': ' + historysep.join(sys.argv) + '\n'
    if 'history' in nc.ncattrs():
        nc.history = historystr + nc.history  # prepend to history string
    else:
        nc.history = historystr

elif solver in cholesky_direct_solvers:
            self.ksp = PETSc.KSP()
            self.ksp.create(PETSc.COMM_WORLD)
            self.ksp.getPC().setType('cholesky')
            self.ksp.getPC().setFactorSolverPackage(solver)
            self.ksp.setType('preonly')

        elif solver in preconditioners:
            self.ksp = PETSc.KSP()
            self.ksp.create(PETSc.COMM_WORLD)
            self.ksp.getPC().setType(solver)
            self.ksp.setType('preonly')

        elif solver in iterative_solvers:
            self.ksp = PETSc.KSP()
            self.ksp.create(PETSc.COMM_WORLD)
            self.ksp.getPC().setType(preconditioner)
            self.ksp.setType(solver)

        self.ksp.setTolerances(self.settings['atol'], self.settings['rtol'],
                               self.settings['maxiter'])