Source code for openpnm.solvers._petsc

import logging
import sys
import numpy as np
import scipy as sp
import scipy.sparse
from openpnm.solvers import IterativeSolver
logger = logging.getLogger(__name__)
try:
    import petsc4py
    # Next line must be before importing PETSc
    petsc4py.init(sys.argv)
    from petsc4py import PETSc
except ModuleNotFoundError:
    pass


__all__ = ['PETScLinearSolver']



[docs]
class PETScLinearSolver(IterativeSolver):
    r"""
    Solves the sparse linear system Ax = b using petsc solvers.

    Notes
    -----
    Parallel computing is supported and matrix partitioning over the
    available cores is automatically handled by running:

    .. code::

        $ mpirun -np num_cores python script.py

    where ``num_cores`` must be substituted with the number of cores.

    """
    def _create_solver(self):
        r"""
        This method creates the petsc sparse linear solver.
        """
        # https://petsc.org/release/docs/manualpages/KSP/KSPType.html
        iterative = [
            'richardson', 'chebyshev', 'cg', 'groppcg', 'pipecg', 'pipecgrr',
            'cgne', 'nash', 'stcg', 'gltr', 'fcg', 'pipefcg', 'gmres',
            'pipefgmres', 'fgmres', 'lgmres', 'dgmres', 'pgmres', 'tcqmr',
            'bcgs', 'ibcgs', 'fbcgs', 'fbcgsr', 'bcgsl', 'pipebcgs', 'cgs',
            'tfqmr', 'cr', 'pipecr', 'lsqr', 'preonly', 'qcg', 'bicg',
            'minres', 'symmlq', 'lcd', 'python', 'gcr', 'pipegcr', 'tsirm',
            'cgls', 'fetidp']
        # https://petsc.org/release/docs/manualpages/PC/PCType.html
        preconditioners = [
            'none', 'jacobi', 'sor', 'lu', 'shell', 'bjacobi', 'mg',
            'eisenstat', 'ilu', 'icc', 'asm', 'gasm', 'ksp', 'composite',
            'redundant', 'spai', 'nn', 'cholesky', 'pbjacobi', 'mat', 'hypre',
            'parms', 'fieldsplit', 'tfs', 'ml', 'galerkin', 'exotic', 'cp',
            'bfbt', 'lsc', 'python', 'pfmg', 'syspfmg', 'redistribute', 'svd',
            'gamg', 'sacusp', 'sacusppoly', 'bicgstabcusp', 'ainvcusp',
            'chowiluviennacl', 'rowscalingviennacl', 'saviennacl', 'bddc',
            'kaczmarz', 'telescope']
        direct_lu = ['mumps', 'superlu_dist', 'umfpack', 'klu']
        direct_cholesky = ['mumps', 'cholmod']
        valid_solvers = iterative + direct_lu + direct_cholesky

        solver = self.solver_type
        preconditioner = self.preconditioner

        if solver not in valid_solvers:
            raise Exception(f"{solver} solver not availabe, choose another solver")
        if preconditioner not in preconditioners:
            raise Exception(f"{preconditioner} not found, choose another preconditioner")

        self.ksp = PETSc.KSP()
        self.ksp.create(PETSc.COMM_WORLD)

        if solver in direct_lu:
            self.ksp.getPC().setType('lu')
            self.ksp.getPC().setFactorSolverType(solver)
            self.ksp.setType('preonly')
        elif solver in direct_cholesky:
            self.ksp.getPC().setType('cholesky')
            self.ksp.getPC().setFactorSolverType(solver)
            self.ksp.setType('preonly')
        elif solver in preconditioners:
            self.ksp.getPC().setType(solver)
            self.ksp.setType('preonly')
        elif solver in iterative:
            self.ksp.getPC().setType(preconditioner)
            self.ksp.setType(solver)

    def _set_tolerances(self, atol=None, rtol=None, maxiter=None):
        r"""
        Set absolute and relative tolerances, and maximum number of iterations.
        """
        atol = self.atol if atol is None else atol
        rtol = self.rtol if rtol is None else rtol
        maxiter = self.maxiter if maxiter is None else maxiter
        # BUG: PETSc misses rtol requirement by ~10-20X -> Report to petsc4py
        self.ksp.setTolerances(atol=None, rtol=rtol/50, max_it=maxiter)

    def _assemble_A(self):
        r"""
        This method creates the petsc sparse coefficients matrix from the
        OpenPNM scipy one. The method also equally decomposes the matrix at
        certain rows into different blocks (each block contains all the
        columns) and distributes them over the pre-assigned cores for parallel
        computing. The method can be used in serial.
        """

        # Create a petsc sparse matrix
        self.petsc_A = PETSc.Mat()
        self.petsc_A.create(PETSc.COMM_WORLD)
        self.petsc_A.setSizes([self.m, self.n])
        self.petsc_A.setType('aij')  # sparse
        self.petsc_A.setUp()

        # Loop over owned block of rows on this processor
        # and insert entry values (for parallel computing).
        self.Istart, self.Iend = self.petsc_A.getOwnershipRange()

        # Assign values to the coefficients matrix from the scipy sparse csr
        size_tmp = self.A.shape
        # Row indices
        csr1 = self.A.indptr[self.Istart:self.Iend+1] - self.A.indptr[self.Istart]
        ind1 = self.A.indptr[self.Istart]
        ind2 = self.A.indptr[self.Iend]
        csr2 = self.A.indices[ind1:ind2]  # column indices
        csr3 = self.A.data[ind1:ind2]  # data
        self.petsc_A = PETSc.Mat().createAIJ(size=size_tmp,
                                             csr=(csr1, csr2, csr3))
        # Communicate off-processor values and setup internal data structures
        # for performing parallel operations
        self.petsc_A.assemblyBegin()
        self.petsc_A.assemblyEnd()

    def _assemble_b_and_x(self):
        r"""
        Initialize the solution vector (self.petsc_x), which is a dense
        matrix (1D vector) and defines the rhs vector (self.petsc_b) from
        the existing data.
        """
        # Get vector(s) compatible with the matrix (same parallel layout)
        # passing same communicator as the A matrix

        # Global solution vector (all the local solutions will return to it)
        self.petsc_s = PETSc.Vec()
        self.petsc_s.create(PETSc.COMM_WORLD)
        self.petsc_s.setSizes(self.m)
        self.petsc_s.setFromOptions()

        self.Istart, self.Iend = self.petsc_s.getOwnershipRange()

        self.petsc_x = (self.petsc_s).duplicate()
        self.petsc_b = (self.petsc_s).duplicate()

        # Set the solution vector to zeros or the given initial guess (if any)
        PETSc.Vec.setArray(self.petsc_x, self.x0[self.Istart: self.Iend])

        # Define the petsc rhs vector from the numpy one
        PETSc.Vec.setArray(self.petsc_b, self.b[self.Istart: self.Iend])


[docs]
    def solve(self, A, b, x0=None, solver_type='cg', precondioner='jacobi',
              maxiter=None, atol=None, rtol=None):
        r"""
        Solves and returns the solution to the linear system, Ax = b.

        This method converts the solution vector from a PETSc.Vec
        instance to a numpy array, and finally destroys all the PETSc
        objects to free memory.

        Parameters
        ----------
        A : csr_matrix
            Coefficients matrix in Ax = b
        b : ndarray
            Right-hand-side vector in Ax = b
        solver_type : str, optional
            Default is the iterative solver 'cg' based on the
            Conjugate Gradient method.
        preconditioner: str, optional
            Default is the 'jacobi' preconditioner, i.e., diagonal
            scaling preconditioning. The preconditioner is used with
            iterative solvers. When a direct solver is used, this
            parameter is ignored.
        factorization_type : str, optional
            The factorization type used with the direct solver.
            Default is 'lu'. This parameter is ignored when an
            iterative solver is used.

        Returns
        -------
        ndarray
            The solution to Ax = b

        Notes
        -----
        Certain combinations of iterative solvers and precondioners
        or direct solvers and factorization types are not supported.
        The summary table of the different possibilities
        can be found
        `here <https://petsc.org/main/overview/linear_solve_table>`_

        """
        self.b = b
        self.A = sp.sparse.csr_matrix(A)
        self.m, self.n = self.A.shape
        self.x0 = np.zeros_like(self.b) if x0 is None else x0

        self.solver_type = solver_type
        self.preconditioner = precondioner

        self.atol = self._get_atol(self.b)
        self.rtol = self._get_rtol(self.x0)

        self._assemble_b_and_x()
        self._assemble_A()

        self._create_solver()
        self._set_tolerances(atol=atol, rtol=rtol, maxiter=maxiter)
        self.ksp.setOperators(self.petsc_A)
        self.ksp.setFromOptions()

        # Solve the linear system
        self.ksp.solve(self.petsc_b, self.petsc_x)

        # Gather the solution to all processors
        gather_to_0, self.petsc_s = PETSc.Scatter().toAll(self.petsc_x)
        gather_to_0.scatter(self.petsc_x, self.petsc_s,
                            PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)

        # Convert solution vector from PETSc.Vec instance to a numpy array
        self.solution = PETSc.Vec.getArray(self.petsc_s)

        # Destroy petsc solver, coefficients matrix, rhs, and solution vectors
        PETSc.KSP.destroy(self.ksp)
        PETSc.Mat.destroy(self.petsc_A)
        PETSc.Vec.destroy(self.petsc_b)
        PETSc.Vec.destroy(self.petsc_x)
        PETSc.Vec.destroy(self.petsc_s)

        # FIXME: fetch exit_code somehow from petsc
        exit_code = 0

        return self.solution, exit_code