### ex9.py from slepc4py demos
try: range = xrange
except: pass

import sys, slepc4py
slepc4py.init(sys.argv)

from petsc4py import PETSc
from slepc4py import SLEPc

Print = PETSc.Sys.Print

def Laplacian2D(m, n):
    """
    Builds discretized Laplacian operator in 2 dimensions.
    """
    # Create matrix for 2D Laplacian operator
    A = PETSc.Mat().create()
    A.setSizes([m*n, m*n])
    A.setFromOptions( )
    A.setUp()
    # Fill matrix
    hx = 1.0/(m-1) # x grid spacing
    hy = 1.0/(n-1) # y grid spacing
    diagv = 2.0*hy/hx + 2.0*hx/hy
    offdx = -1.0*hy/hx
    offdy = -1.0*hx/hy
    Istart, Iend = A.getOwnershipRange()
    for I in range(Istart, Iend):
        A[I,I] = diagv
        i = I//n    # map row number to
        j = I - i*n # grid coordinates
        if i> 0  : J = I-n; A[I,J] = offdx
        if i< m-1: J = I+n; A[I,J] = offdx
        if j> 0  : J = I-1; A[I,J] = offdy
        if j< n-1: J = I+1; A[I,J] = offdy
    A.assemble()
    return A

def QuasiDiagonal(N):
    """
    Builds matrix diag(2)+[4 -1; -1 -1]
    """
    # Create matrix
    B = PETSc.Mat().create()
    B.setSizes([N, N])
    B.setFromOptions( )
    B.setUp()
    # Fill matrix
    Istart, Iend = B.getOwnershipRange()
    for I in range(Istart, Iend):
        B[I,I] = 2.0
    if Istart==0:
        B[0,0] = 6.0
        B[0,1] = -1.0
        B[1,0] = -1.0
        B[1,1] = 1.0
    B.assemble()
    return B

### testing with B operator
m = 10
n = 10
A = Laplacian2D(m,n)
B = QuasiDiagonal(m*n)

class BOperator(object):
    def mult(this, mat, x, y):
        y.set(0)
        y.axpy(1.0, B*x)

Bop = PETSc.Mat().create()
Bop.setSizes(*B.size)
Bop.setType(Bop.Type.PYTHON)
Bop.setPythonContext(BOperator())
Bop.setUp()

E = SLEPc.EPS().create()
E.setDimensions(10,PETSc.DECIDE)
E.setType(E.Type.KRYLOVSCHUR)
E.setProblemType(SLEPc.EPS.ProblemType.GHEP)
E.setWhichEigenpairs(E.Which.LARGEST_MAGNITUDE)
E.setTolerances(1e-10)
E.setOperators(A, Bop)
E.setMonitor(None)
E.setFromOptions()

st = E.getST()
ksp = st.getKSP()
ksp.setOperators(Bop,B)
pc = ksp.getPC()
pc.setType("ilu")

#E.view()
E.solve()

its = E.getIterationNumber()
Print("Number of iterations of the method: %i" % its)
sol_type = E.getType()
Print("Solution method: %s" % sol_type)
nev, ncv, mpd = E.getDimensions()
Print("Number of requested eigenvalues: %i" % nev)
tol, maxit = E.getTolerances()
Print("Stopping condition: tol=%.4g, maxit=%d" % (tol, maxit))
nconv = E.getConverged()
Print("Number of converged eigenpairs: %d" % nconv)