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

def __init__(self, vs):
        if vs.enable_cyclic_x:
            boundary_type = ('periodic', 'ghosted')
        else:
            boundary_type = ('ghosted', 'ghosted')

        self._da = PETSc.DMDA().create(
            [vs.nx, vs.ny],
            stencil_width=1,
            stencil_type='star',
            comm=rs.mpi_comm,
            proc_sizes=rs.num_proc,
            boundary_type=boundary_type,
            ownership_ranges=[
                (vs.nx // rs.num_proc[0],) * rs.num_proc[0],
                (vs.ny // rs.num_proc[1],) * rs.num_proc[1],
            ]
        )

        self._matrix, self._boundary_fac = self._assemble_poisson_matrix(vs)

        petsc_options = PETSc.Options()

def __init__(self, vs):
        if vs.enable_cyclic_x:
            boundary_type = ("periodic", "ghosted")
        else:
            boundary_type = ("ghosted", "ghosted")

        self._da = PETSc.DMDA().create(
            [vs.nx, vs.ny],
            stencil_width=1,
            stencil_type="star",
            comm=rst.mpi_comm,
            proc_sizes=rs.num_proc,
            boundary_type=boundary_type,
            ownership_ranges=[
                (vs.nx // rs.num_proc[0],) * rs.num_proc[0],
                (vs.ny // rs.num_proc[1],) * rs.num_proc[1],
            ]
        )

        self._matrix, self._boundary_fac = self._assemble_poisson_matrix(vs)

        petsc_options = PETSc.Options()

space_size = space_comm.Get_size()

    if len(sys.argv) == 2:
        color = int(world_rank % int(sys.argv[1]))
    else:
        color = int(world_rank / world_size)

    time_comm = comm.Split(color=color)
    time_rank = time_comm.Get_rank()
    time_size = time_comm.Get_size()

    print("IDs (world, space, time):  %i / %i -- %i / %i -- %i / %i" % (world_rank, world_size, space_rank, space_size,
                                                                        time_rank, time_size))

    n = 7
    da = PETSc.DMDA().create([n, n], stencil_width=1, comm=space_comm)

    x = da.createGlobalVec()
    xa = da.getVecArray(x)
    (xs, xe), (ys, ye) = da.getRanges()
    for i in range(xs, xe):
        for j in range(ys, ye):
            xa[i, j] = np.sin(2 * np.pi * (i + 1) / (n + 1)) * np.sin(2 * np.pi * (j + 1) / (n + 1))
    print('x=', x.getArray())
    print('x:', x.getSizes(), da.getRanges())
    print()

    if time_rank == 0:
        print('send', time_rank)
        time_comm.send(x.getArray(), dest=1, tag=0)
    else:
        print('recv', time_rank)

Initialization routine

        Args:
            init: can either be a tuple (one int per dimension) or a number (if only one dimension is requested)
                  or another DMDA object
            val: initial value (default: None)
        Raises:
            DataError: if init is none of the types above
        """

        # if init is another petsc data type, do a copy (init by copy)
        if isinstance(init, type(self)):
            self.comp1 = petsc_data(init.comp1)
            self.comp2 = petsc_data(init.comp2)
        # if init is a DMDA, create an empty object
        elif isinstance(init, PETSc.DMDA):
            self.comp1 = petsc_data(init)
            self.comp2 = petsc_data(init)
        # something is wrong, if none of the ones above hit
        else:
            raise DataError('something went wrong during %s initialization' % type(self))

# define the Dirichlet boundary
        if 'comm' not in problem_params:
            problem_params['comm'] = PETSc.COMM_WORLD
        if 'sol_tol' not in problem_params:
            problem_params['sol_tol'] = 1E-10
        if 'sol_maxiter' not in problem_params:
            problem_params['sol_maxiter'] = None

        # these parameters will be used later, so assert their existence
        essential_keys = ['nvars', 'lambda0', 'nu', 'interval']
        for key in essential_keys:
            if key not in problem_params:
                msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys()))
                raise ParameterError(msg)
        # create DMDA object which will be used for all grid operations (boundary_type=3 -> periodic BC)
        da = PETSc.DMDA().create([problem_params['nvars']], dof=1, stencil_width=1, comm=problem_params['comm'])

        # invoke super init, passing number of dofs, dtype_u and dtype_f
        super(petsc_fisher, self).__init__(init=da, dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params)

        # compute dx, dy and get local ranges
        self.dx = (self.params.interval[1] - self.params.interval[0]) / (self.params.nvars - 1)
        # print(self.init.getRanges())
        (self.xs, self.xe) = self.init.getRanges()[0]

        # compute discretization matrix A and identity
        self.A = self.__get_A()
        self.localX = self.init.createLocalVec()

        self.ksp_itercount = 0
        self.ksp_ncalls = 0

# define the Dirichlet boundary
        if 'comm' not in problem_params:
            problem_params['comm'] = PETSc.COMM_WORLD
        if 'sol_tol' not in problem_params:
            problem_params['sol_tol'] = 1E-10
        if 'sol_maxiter' not in problem_params:
            problem_params['sol_maxiter'] = None

        # these parameters will be used later, so assert their existence
        essential_keys = ['nvars', 'Du', 'Dv', 'A', 'B']
        for key in essential_keys:
            if key not in problem_params:
                msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys()))
                raise ParameterError(msg)
        # create DMDA object which will be used for all grid operations (boundary_type=3 -> periodic BC)
        da = PETSc.DMDA().create([problem_params['nvars'][0], problem_params['nvars'][1]], dof=2, boundary_type=3,
                                 stencil_width=1, comm=problem_params['comm'])

        # invoke super init, passing number of dofs, dtype_u and dtype_f
        super(petsc_grayscott, self).__init__(init=da, dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params)

        # compute dx, dy and get local ranges
        self.dx = 100.0 / (self.params.nvars[0])
        self.dy = 100.0 / (self.params.nvars[1])
        (self.xs, self.xe), (self.ys, self.ye) = self.init.getRanges()

        # compute discretization matrix A and identity
        self.A = self.__get_A()
        self.localX = self.init.createLocalVec()

        # setup nonlinear solver
        self.snes = PETSc.SNES()

# if j < ny-1: u_n = x[i, j+1] # north
                # u_xx = (u_e - 2*u + u_w)*hy/hx
                # u_yy = (u_n - 2*u + u_s)*hx/hy
                # y[i, j] = u_xx + u_yy
        mat.assemble()
        mat.mult(x, y)
        # if J != P: J.assemble()  # matrix-free operator
        return PETSc.Mat.Structure.SAME_NONZERO_PATTERN

OptDB = PETSc.Options()

n  = OptDB.getInt('n', 4)
nx = OptDB.getInt('nx', n)
ny = OptDB.getInt('ny', n)

da = PETSc.DMDA().create([nx, ny], stencil_width=1)
pde = Poisson2D(da)

x = da.createGlobalVec()
b = da.createGlobalVec()
# A = da.createMat('python')
A = PETSc.Mat().createPython(
    [x.getSizes(), b.getSizes()], comm=da.comm)
A.setPythonContext(pde)
A.setUp()

ksp = PETSc.KSP().create()
ksp.setOperators(A)
ksp.setType('cg')
pc = ksp.getPC()
pc.setType('none')
ksp.setFromOptions()

def main():

    n = 4
    da = PETSc.DMDA().create([n, n], stencil_width=1)

    rank = PETSc.COMM_WORLD.getRank()

    x = da.createGlobalVec()
    xa = da.getVecArray(x)
    (xs, xe), (ys, ye) = da.getRanges()
    for i in range(xs, xe):
        for j in range(ys, ye):
            xa[i, j] = j*n + i
    print('x=', rank, x.getArray(), xs, xe, ys, ye)

    A = da.createMatrix()
    A.setType('aij')  # sparse
    A.setFromOptions()

    Istart, Iend = A.getOwnershipRange()

problem_params['sol_tol'] = 1E-10
        if 'sol_maxiter' not in problem_params:
            problem_params['sol_maxiter'] = None

        essential_keys = ['cnvars', 'nu', 'freq', 'comm', 'refine']
        for key in essential_keys:
            if key not in problem_params:
                msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys()))
                raise ParameterError(msg)

        # make sure parameters have the correct form
        if len(problem_params['cnvars']) != 2:
            raise ProblemError('this is a 2d example, got %s' % problem_params['cnvars'])

        # create DMDA object which will be used for all grid operations
        da = PETSc.DMDA().create([problem_params['cnvars'][0], problem_params['cnvars'][1]], stencil_width=1,
                                 comm=problem_params['comm'])
        for _ in range(problem_params['refine']):
            da = da.refine()

        # invoke super init, passing number of dofs, dtype_u and dtype_f
        super(heat2d_petsc_forced, self).__init__(init=da, dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params)

        # compute dx, dy and get local ranges
        self.dx = 1.0 / (self.init.getSizes()[0] - 1)
        self.dy = 1.0 / (self.init.getSizes()[1] - 1)
        (self.xs, self.xe), (self.ys, self.ye) = self.init.getRanges()

        # compute discretization matrix A and identity
        self.A = self.__get_A()
        self.Id = self.__get_Id()

u_e = u_w = u_n = u_s = 0
                if i > 0:    u_w = x[i-1, j] # west
                if i < mx-1: u_e = x[i+1, j] # east
                if j > 0:    u_s = x[i, j-1] # south
                if j < my-1: u_n = x[i, j+1] # north
                u_xx = (u_e - 2*u + u_w) / hx ** 2
                u_yy = (u_n - 2*u + u_s) / hy ** 2
                y[i, j] = u_xx + u_yy

OptDB = PETSc.Options()

n  = OptDB.getInt('n', 4)
nx = OptDB.getInt('nx', n)
ny = OptDB.getInt('ny', n)

da = PETSc.DMDA().create([nx, ny], stencil_width=1)
pde = Poisson2D(da)

x = da.createGlobalVec()
b = da.createGlobalVec()
# A = da.createMat('python')
A = PETSc.Mat().createPython(
    [x.getSizes(), b.getSizes()], comm=da.comm)
print(A.getSize())
A.setPythonContext(pde)
A.setUp()

y = da.createGlobalVec()

pde.formRHS(x, val=1.0)
A.mult(x, b)
pde.formRHS(y, val=-2.0 * (2.0 * np.pi) ** 2)