### ====================================================================
###  Python-file
###     author:          Ethan T. Coon
###     filename:        sp_solver.py
###     version:         
###     created:         08 June 2009
###       on:            16:17:44 MDT
###     last modified:   07 February 2011
###       at:            09:24:24 MST
###     URL:             http://www.ldeo.columbia.edu/~ecoon/
###     email:           etc2103 _at_ columbia.edu
###  
### ====================================================================

"""Saddle point problems solver:

Matrix-free in PETSc

Example of the versatility of using PETSc and petsc4py's matrix-free
contexts.  Solves the block system:

    [ A   B ] u       v
    |       |       =    
    [ B^T C ] lmbda   mu

using, as a left preconditioner:

[ A'  0 ]
[ 0   S ]

where A' is a preconditioner for A and S = C - B^T*A^-1*B is the Schur 
Complement of the saddle point problem.


If run as main, this solves 1D Laplace's equation, using Lagrange
Multipliers to do the boundary conditions.  To do this, call:

python lm_solver.py -schur_ksp_monitor -outer_ksp_monitor -n NUM_POINTS

which shows the iterations of the 'outer' ksp and the 'inner' ksp which
inverts the schur complement.  Apologies for crufty formatting, as the
name doesn't seem to show up in the iterations in PETSc, so the various 
monitors get interleaved...
"""

import sys
import numpy


def get_python_matrix( shell ):
    M = PETSc.Mat().create()
    M.setSizes(shell.size)
    M.setType('python')
    M.setPythonContext(shell)
    return M
                     

class SchurMatShell(object):
    """Matrix shell for the Schur complement, C - B^T*A^-1*B"""
    
    def __init__( self, context ):
        self.context = context
        size = self.context.B.size[1]
        self.size = (size, size)

        self.u = context.u.duplicate()
        self.v = context.u.duplicate()

    def create( self, M ):
        pass

    def mult( self, M, x, y ):
        """All matrix-free contexts must provide the action of the
        matrix

        y <-- M*x
        """

        ctx = self.context
        u,v = self.u, self.v

        ctx.B.mult(x,u)
        v.zeroEntries()

        # two options -- we can either solve the linear system
        #   to get A^-1, or simply apply the approximate inverse
        #   of A, i.e. use A's preconditioner.  Likely the best
        #   choice is a low-tolerance solve of A, (i.e. just a few
        #   iterations of CG)

        #ctx.A_ksp.solve(u,v) # solve a ksp to get the action of A^-1
        ctx.A_pc.apply(u,v) # use A's preconditioner as an approximate
                            # action of A^-1
        ctx.B.multTranspose(v,y)
        y.scale(-1.)

        if ctx.C is not None:
            ctx.C.multAdd(x, y, y)
        return

class SchurPCShell(object):
    def __init__( self, context ):
        self.context = context

        # Schur Complement
        self.S = get_python_matrix(SchurMatShell(context))

        ksp = PETSc.KSP().create()
        ksp.setName('Schur solver')
        ksp.setType('gmres')
        ksp.setOptionsPrefix('schur_')
        ksp.setOperators(self.S,self.S)
        pc = ksp.getPC()
        pc.setType('none')                          
        ksp.setFromOptions()
        self.ksp = ksp

        self.rhs = context.lmbda.duplicate()

    def apply( self, pc, x, y ):
        """All pc contexts must provide the action of the inverse.

        y <-- S^-1*x
        """
        
        # to apply the Schur complement as a preconditioner, we have
        # to solve the system.  This should also be a lightweight solve?
        rhs = self.rhs
        x.copy(rhs)
        self.ksp.solve(rhs,y)
        try:
            assert self.ksp.getConvergedReason() > 0
        except AssertionError:
            raise ValueError('Schur KSP diverged with reason: %d'%self.ksp.getConvergedReason())
        return


class SaddlePointMatShell(object):
    """The block matrix: Matrix free shell for the full system"""
    
    def __init__( self, context ):
        self.context = context
        self.size = context.size

        self.u = context.u.duplicate()
        self.v = context.u.duplicate()
        self.lmbda = context.lmbda.duplicate()
        self.mu = context.lmbda.duplicate()

    def create( self, M ):
        pass

    def mult( self, M, x, y ):
        """y <- M * x"""
        ctx = self.context
        A,B,C = ctx.A, ctx.B, ctx.C

        u, v = self.u, self.v
        lmbda, mu = self.lmbda, self.mu

        # pull blocks out of the global vector
        u[...] = x.getValues(ctx.u_is)
        lmbda[...] = x.getValues(ctx.l_is)

        # do the block multiplication
        #  - upper row
        A.mult(u, v)
        B.multAdd(lmbda, v, v)
        #  - lower row
        B.multTranspose(u, mu)
        if C is not None:
            C.multAdd(lmbda, mu, mu)

        # reinsert the result
        y.setValues(ctx.u_is, v[...])
        y.setValues(ctx.l_is, mu[...])
        y.assemble()
        return

class SaddlePointPCShell(object):
    """The block preconditioner:  diag( A, S )^-1"""
    
    def __init__( self, context ):
        self.context = context

        # Schur Complement Inversion
        S_pc = PETSc.PC().create()
        S_pc.setType('python')
        S_pc.setPythonContext(SchurPCShell(context))
        S = S_pc.getPythonContext().S
        S_pc.setOperators(S, S)
        self.S_pc = S_pc

        self.u = context.u.duplicate()
        self.v = context.u.duplicate()
        self.lmbda = context.lmbda.duplicate()
        self.mu = context.lmbda.duplicate()

    def apply( self, pc, x, y ):
        """y <-- M^-1 * x"""
        ctx = self.context
        u, v = self.u, self.v
        lmbda, mu = self.lmbda, self.mu

        # upper row, preconditioned by A
        u[...] = x.getValues(ctx.u_is)
        ctx.A_pc.apply(u,v)
        y.setValues(ctx.u_is, v[...])

        # lower row, preconditioned by S
        lmbda[...] = x.getValues(ctx.l_is)
        self.S_pc.apply(lmbda, mu)
        y.setValues(ctx.l_is, mu[...])
        y.assemble()
        return 

class LagrangeMultiplierContext(object):
    """Lagrange Multiplier MF context

    [ A   B ]u       v
    |       |      =    
    [ B^T C ]lmbda   mu

    where A of size  [n, n]
    and C of size [m, m]

    Just an empty struct of data to pass around.
    """
    def __init__( self, A, B, C=None, A_pc=None ):
        self.A = A
        self.B = B
        self.C = C
        if A_pc is None:
            A_pc = A

        n,m = A.size[1], B.size[1]
        self.size = (n+m, n+m)

        # vectors and index sets
        self.lmbda, self.u = B.getVecs()
         
        self.u_is = PETSc.IS().createGeneral(range(n)) # index set for the A-block
        self.l_is = PETSc.IS().createGeneral(range(n,n+m)) # index set for the S-block

        ksp = PETSc.KSP().create()
        ksp.setName('A_solver')
        ksp.setOptionsPrefix('A_')
        ksp.setType('cg')
        ksp.setOperators(A, A_pc)
        pc = ksp.getPC()
        pc.setType('icc')                          
        pc.setFromOptions()
        ksp.setFromOptions()
        self.A_ksp = ksp
        self.A_pc = pc

        
if __name__ == '__main__':
    import petsc4py, sys
    petsc4py.init(sys.argv)
    from petsc4py import PETSc

    opts = PETSc.Options()
    n = opts.getInt('n', 8)

    # submatrix A (1D Laplacian)
    A = PETSc.Mat().createAIJ((n,n))
    v = PETSc.Vec().createSeq(n)
    diag = [2 for i in range(n-2)]
    diag.insert(0,1)
    diag.append(1)
    v[...] = diag
    
    A.setDiagonal(v)
    A[0,1] = -1.
    for lcv in range(1,n-1):
        A[lcv, [lcv-1, lcv+1]] = [-1.,-1.]
    A[n-1,n-2] = -1.
    A.assemble()

    # submatrix B (integral of phi*lambda on boundary)
    B = PETSc.Mat().createAIJ((n,2))
    B[0,0] = -1. # du/dx(0) = 1
    B[n-1,1] = -1. # du/dx(1) = 1
    B.assemble()
    
    # user context
    ctx = LagrangeMultiplierContext(A,B)    
    
    M_shell = SaddlePointMatShell(ctx)
    M = get_python_matrix(M_shell)

    # rhs, soln
    x,b = M.getVecs()
    x.set(0.)
    b.set(0.)
    b[n+1] = -1.
    b.assemble()

    # solver -- GMRES with block preconditioner
    ksp = PETSc.KSP().create()
    ksp.setName('Saddle Point Solver')
    ksp.setOptionsPrefix('outer_')
    pc = ksp.getPC()
    ksp.setType('gmres')
    pc.setType('python')
    pc.setPythonContext(SaddlePointPCShell(ctx))
    ksp.setOperators(M, M)
    ksp.setFromOptions()

    # solve:
    x.zeroEntries()
    ksp.solve(b,x)
    print 'solved %s with converged reason %d'%(ksp.name, ksp.getConvergedReason())

    # plot
    from matplotlib import pyplot as plt
    xs = numpy.arange(0,1.00001,1./(n-1))
    plt.plot(xs, x.getValues(ctx.u_is), 'b-o')
    plt.show()
    
else:
    from petsc4py import PETSc