# Poisson in 2D
# =============
#
# Solve a constant coefficient Poisson problem on a regular grid. The source
# code for this demo can be `downloaded here <../../_static/poisson2d.py>`__
#
# .. math::
#
#     - u_{xx} - u_{yy} = 1 \quad\textsf{in}\quad [0,1]^2\\
#     u = 0 \quad\textsf{on the boundary.}

# This is a naïve, parallel implementation, using :math:`n` interior grid
# points per dimension and a lexicographic ordering of the nodes.

# This code is kept as simple as possible. However, simplicity comes at a
# price. Here we use a naive decomposition that does not lead to an optimal
# communication complexity for the matrix-vector product. An optimal complexity
# decomposition of a structured grid could be achieved using `PETSc.DMDA`.
#
# This demo is structured as a script to be executed using:
#
# .. code-block:: console
#
#   $ python poisson2d.py
#
# potentially with additional options passed at the end of the command.
#
# At the start of your script, call `petsc4py.init` passing `sys.argv` so that
# command-line arguments to the script are passed through to PETSc.

import sys
import petsc4py

petsc4py.init(sys.argv)

# The full PETSc4py API is to be found in the `petsc4py.PETSc` module.

from petsc4py import PETSc

# PETSc is extensively programmable using the `PETSc.Options` database. For
# more information see `working with PETSc Options <petsc_options>`.

OptDB = PETSc.Options()

# Grid size and spacing using a default value of ``5``. The user can specify a
# different number of points in each direction by passing the ``-n`` option to
# the script.

n = OptDB.getInt('n', 5)
h = 1.0 / (n + 1)

# Matrices are instances of the `PETSc.Mat` class.

A = PETSc.Mat()

# Create the underlying PETSc C Mat object.
# You can omit the ``comm`` argument if your objects live on
# `PETSc.COMM_WORLD` but it is a dangerous choice to rely on default values
# for such important arguments.

A.create(comm=PETSc.COMM_WORLD)

# Specify global matrix shape with a tuple.

A.setSizes((n * n, n * n))

# The call above implicitly assumes that we leave the parallel decomposition of
# the matrix rows to PETSc by using `PETSc.DECIDE` for local sizes.
# It is equivalent to:
#
# .. code-block:: python
#
#     A.setSizes(((PETSc.DECIDE, n * n), (PETSc.DECIDE, n * n)))

# Here we use a sparse matrix of AIJ type
# Various `matrix formats <petsc4py.PETSc.Mat.Type>` can be selected:

A.setType(PETSc.Mat.Type.AIJ)

# Finally we allow the user to set any options they want to on the matrix from
# the command line:

A.setFromOptions()

# Insertion into some matrix types is vastly more efficient if we preallocate
# space rather than allow this to happen dynamically. Here we hint the number
# of nonzeros to be expected on each row.

A.setPreallocationNNZ(5)

# We can now write out our finite difference matrix assembly using conventional
# Python syntax. `Mat.getOwnershipRange` is used to retrieve the range of rows
# local to this processor.


def index_to_grid(r):
    """Convert a row number into a grid point."""
    return (r // n, r % n)


rstart, rend = A.getOwnershipRange()
for row in range(rstart, rend):
    i, j = index_to_grid(row)
    A[row, row] = 4.0 / h**2
    if i > 0:
        column = row - n
        A[row, column] = -1.0 / h**2
    if i < n - 1:
        column = row + n
        A[row, column] = -1.0 / h**2
    if j > 0:
        column = row - 1
        A[row, column] = -1.0 / h**2
    if j < n - 1:
        column = row + 1
        A[row, column] = -1.0 / h**2

# At this stage, any exchange of information required in the matrix assembly
# process has not occurred. We achieve this by calling `Mat.assemblyBegin` and
# then `Mat.assemblyEnd`.

A.assemblyBegin()
A.assemblyEnd()

# We set up an additional option so that the user can print the matrix by
# passing ``-view_mat`` to the script.

A.viewFromOptions('-view_mat')

# PETSc represents all linear solvers as preconditioned Krylov subspace methods
# of type `PETSc.KSP`. Here we create a KSP object for a conjugate gradient
# solver preconditioned with an algebraic multigrid method.

ksp = PETSc.KSP()
ksp.create(comm=A.getComm())
ksp.setType(PETSc.KSP.Type.CG)
ksp.getPC().setType(PETSc.PC.Type.GAMG)

# We set the matrix in our linear solver and allow the user to program the
# solver with options.

ksp.setOperators(A)
ksp.setFromOptions()

# Since the matrix knows its size and parallel distribution, we can retrieve
# appropriately-scaled vectors using `Mat.createVecs`. PETSc vectors are
# objects of type `PETSc.Vec`. Here we set the right-hand side of our system to
# a vector of ones, and then solve.

x, b = A.createVecs()
b.set(1.0)
ksp.solve(b, x)

# Finally, allow the user to print the solution by passing ``-view_sol`` to the
# script.

x.viewFromOptions('-view_sol')

# Things to try
# -------------
#
# - Show the solution with ``-view_sol``.
# - Show the matrix with ``-view_mat``.
# - Change the resolution with ``-n``.
# - Use a direct solver by passing ``-ksp_type preonly -pc_type lu``.
# - Run in parallel on two processors using:
#
#   .. code-block:: console
#
#       mpiexec -n 2 python poisson2d.py