How to use petsc4py - 10 common examples

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()

        # setup krylov method
        self._ksp = PETSc.KSP()
        self._ksp.create(rs.mpi_comm)
        self._ksp.setOperators(self._matrix)
        self._ksp.setType('bcgs')
        self._ksp.setTolerances(atol=0, rtol=vs.congr_epsilon, max_it=vs.congr_max_iterations)

        # preconditioner
        self._ksp.getPC().setType('hypre')
        petsc_options['pc_hypre_type'] = 'boomeramg'
        self._ksp.getPC().setFromOptions()

        self._rhs_petsc = self._da.createGlobalVec()
        self._sol_petsc = self._da.createGlobalVec()

# 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))

north_diag = vs.hur[2:-2, 3:-1] / vs.dyu[np.newaxis, 2:-2] / \
            vs.dyt[np.newaxis, 3:-1] * vs.cost[np.newaxis, 3:-1] / vs.cosu[np.newaxis, 2:-2]
        south_diag = vs.hur[2:-2, 2:-2] / vs.dyu[np.newaxis, 2:-2] / \
            vs.dyt[np.newaxis, 2:-2] * vs.cost[np.newaxis, 2:-2] / vs.cosu[np.newaxis, 2:-2]

        # construct sparse matrix
        cf = tuple(diag for diag in (
            boundary_mask * main_diag + (1 - boundary_mask),
            boundary_mask * east_diag,
            boundary_mask * west_diag,
            boundary_mask * north_diag,
            boundary_mask * south_diag
        ))

        row = PETSc.Mat.Stencil()
        col = PETSc.Mat.Stencil()

        ij_offsets = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)]

        (i0, i1), (j0, j1) = self._da.getRanges()
        for j in range(j0, j1):
            for i in range(i0, i1):
                iloc, jloc = i % (vs.nx // rs.num_proc[0]), j % (vs.ny // rs.num_proc[1])
                row.index = (i, j)

                for diag, offset in zip(cf, ij_offsets):
                    io, jo = (i + offset[0], j + offset[1])
                    col.index = (io, jo)
                    matrix.setValueStencil(row, col, diag[iloc, jloc])

        matrix.assemble()

vs.dxt[2:-2, np.newaxis] / vs.cosu[np.newaxis, 2:-2]**2
        north_diag = vs.hur[2:-2, 3:-1] / vs.dyu[np.newaxis, 2:-2] / \
            vs.dyt[np.newaxis, 3:-1] * vs.cost[np.newaxis, 3:-1] / vs.cosu[np.newaxis, 2:-2]
        south_diag = vs.hur[2:-2, 2:-2] / vs.dyu[np.newaxis, 2:-2] / \
            vs.dyt[np.newaxis, 2:-2] * vs.cost[np.newaxis, 2:-2] / vs.cosu[np.newaxis, 2:-2]

        # construct sparse matrix
        cf = tuple(diag for diag in (
            boundary_mask * main_diag + (1 - boundary_mask),
            boundary_mask * east_diag,
            boundary_mask * west_diag,
            boundary_mask * north_diag,
            boundary_mask * south_diag
        ))

        row = PETSc.Mat.Stencil()
        col = PETSc.Mat.Stencil()

        ij_offsets = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)]

        (i0, i1), (j0, j1) = self._da.getRanges()
        for j in range(j0, j1):
            for i in range(i0, i1):
                iloc, jloc = i % (vs.nx // rs.num_proc[0]), j % (vs.ny // rs.num_proc[1])
                row.index = (i, j)

                for diag, offset in zip(cf, ij_offsets):
                    io, jo = (i + offset[0], j + offset[1])
                    col.index = (io, jo)
                    matrix.setValueStencil(row, col, diag[iloc, jloc])

        matrix.assemble()

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()

def _setPETSc(self,petsc_file):
        """The function takes a file with petsc options and sets the options globally.

        petsc_file : str
            string with the location of the file
        """
        # First, clear any existing PETSc options settings
        for key in petsc4py.PETSc.Options().getAll():
            petsc4py.PETSc.Options().delValue(key)
        # Second, collect and add new PETSc options
        petsc_options = []
        with open(petsc_file) as petsc_file:
            data = petsc_file.readlines()
        def strip_comments(line):
            if '#' in line:
                line = line[:line.index('#')]
            return line
        stripped_data = [strip_comments(line) for line in data]
        petsc_options = ''.join(stripped_data).split('\n')
        new_petsc = []
        for item in petsc_options:
            if item != '':
                new = item.split()
                new[0] = new[0][1:]

def test_ISRead(self):
        """Test reading an IS"""
        indices = np.array([3,4,5])
        anis = PETSc.IS().createGeneral(list(indices))
        viewer = PETSc.Viewer().createBinary('test.dat', PETSc.Viewer.Mode.W)
        anis.view(viewer)
        viewer.destroy()
        anis.destroy()

        result, = PetscBinaryIO().readBinaryFile('test.dat')
        self.assertTrue((indices == result).all())

u2.create(grid, "", PISM.WITH_GHOSTS)
        u2.copy_from(ssarun.ssa.solution())

        Td_fd = PISM.IceModelVec2V()
        Td_fd.create(grid, "", PISM.WITH_GHOSTS)
        Td_fd.copy_from(u2)
        Td_fd.add(-1, u1)
        Td_fd.scale(1. / eps)

        d_Td = PISM.IceModelVec2V()
        d_Td.create(grid, "", PISM.WITH_GHOSTS)
        d_Td.copy_from(Td_fd)
        d_Td.add(-1, Td)

        n_Td_fd = Td_fd.norm(PETSc.NormType.NORM_2)
        n_Td_l2 = Td.norm(PETSc.NormType.NORM_2)
        n_Td_l1 = Td.norm(PETSc.NormType.NORM_1)
        n_Td_linf = Td.norm(PETSc.NormType.NORM_INFINITY)

        n_d_Td_l2 = d_Td.norm(PETSc.NormType.NORM_2)
        n_d_Td_l1 = d_Td.norm(PETSc.NormType.NORM_1)
        n_d_Td_linf = d_Td.norm(PETSc.NormType.NORM_INFINITY)

        PISM.verbPrintf(1, grid.com, "(i,j)=(%d,%d)\n" % (i, j))
        PISM.verbPrintf(1, grid.com, "apply_linearization(d): l2 norm %.10g; finite difference %.10g\n" %
                        (n_Td_l2, n_Td_fd))

        r_d_l2 = 0
        if n_Td_l2 != 0:
            r_d_l2 = n_d_Td_l2 / n_Td_l2
        r_d_l1 = 0
        if n_Td_l1 != 0: