Expose KSPSetComputeRHS and KSPSetComputeOperators in petsc4py

[Francesco Caimmi]

It would be nice to have access to KSPSetComputeRHS and KSPSetComputeOperators functions in petsc4py.

I open this issue to comply with Lisandro's request as per http://lists.mcs.anl.gov/pipermail/petsc-users/2015-July/026256.html.

Thanks.

[Lisandro Dalcin]

Fixed in branch maint (0f8fc08).

@francesco_caimmi Could you please give it a try and confirm it is working as expected?

[Francesco Caimmi]

Hi, thanks for exposing the functions. However I wasn't able to reproduce the expected behavior. I tried to translate ksp/tutorials/ex25.c to python but when I call ksp.solve I get a TypeError: iteration over non-sequence error. The code I used is at the end of the comment. However note that, at variance with the C code since ksp.solve cannot accept None as argument (the C function call in ex25 reads KSPSolve(ksp,NULL,NULL);), I had to explicitly create vectors b and x before calling ksp.solve(b,x). I do not know if this might have an influence on the process.If I go with the ksp.setComputeOperators route I get the expected result (more or less) so I guess the code is correct. Note also that is not clear to me if in your implementation the function defining the known vector and the operator are supposed to return an integer (as in the C implementation) or the vector/operator itself.

Here is the full traceback:

File "./ksp_ex25.py", line 179, in <module> main(k,e) File "./ksp_ex25.py", line 155, in main ksp.solve(b,x) File "PETSc/KSP.pyx", line 377, in petsc4py.PETSc.KSP.solve (src/petsc4py.PETSc.c:150958) File "PETSc/petscksp.pxi", line 241, in petsc4py.PETSc.KSP_ComputeRHS (src/petsc4py.PETSc.c:31355) TypeError: iteration over non-sequence

and here is the code

import sys
import petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
import numpy as np

class AppCtx():
    """
    Conteiner class to store user data

    Parameters
    ------------
    k: int
        sine arg coefficient
    e: float
        sine multiplying factor

    """
    def __init__(self, k, e):
        self.k = k
        self.e = e

def computeMatrix(ksp, J, jac, Ctx):
    """
    Computes the system matrix
    Parameters
    -----------
    ksp: PETSc ksp object
        the ksp object

    J: PETSc Mat
        the linear operator

    jac:PETSc Mat
         preconditioning matrix

    Ctx: AppCtx istance
        teh user context

    Returns
    ---------
    0
    """
    da = ksp.getDM()
    mx = da.getSizes()[0]
    xs , xm = da.getCorners()
#    row = PETSc.Mat.Stencil()
#    col = PETSc.Mat.Stencil()
    xs = xs[0]
    xm = xm[0]
    e = Ctx.e
    k = Ctx.k
    h = 1.0/(mx-1)
    values = np.zeros(3)
    #I use setValue rather than stencils cause I cannot figure out how
    #they work in petsc4py
    for i in range(xs,xs+xm):
        if i==0 or i==mx-1:
            values[0]=2.0
            jac.setValues(i,i,values[0])
        else:
            xlow = h*i-0.5*h
            xhigh = xlow+h
            values[0] = (-1.0-e*np.sin(2*np.pi*k*xlow))/h
            values[1] = (2.0+e*np.sin(2*np.pi*k*xlow)+e*np.sin(2*np.pi*k*xhigh))/h
            values[2] = (-1.0-e*np.sin(2*np.pi*k*xhigh))/h
            jac.setValues(i,[0,1,2],values)

    jac.assemblyBegin()
    jac.assemblyEnd()
    return 0

def computeRHS(ksp,b, Ctx):
    """

    Parameters
    -----------
    ksp: PETSc ksp object
        the ksp object

    b: PETSc Vec
        a vector

    Ctx: AppCtx istance
        teh user context

    Return
    --------
    0
    """

    da = ksp.getDM()
    #there is no DMDAGetInfo in petsc4py, so use getSizes
    mx = da.getSizes()[0]
    h = 1.0/(mx-1)
    b.set(h)
    idx = [0,mx-1]
    values = [0.0,0.0]
    b.setValues(idx,values)
    b.assemblyBegin()
    b.assemblyEnd()
    return 0

def main(k,e):
    """
    Parameters
    ------------
    k: int
        sine arg coefficient
    e: float
        sine multiplying factor

    """
    #init user context
    Ctx = AppCtx(k,e)
    #set some utilities
    comm = PETSc.COMM_WORLD
    #create the ksp context
    ksp = PETSc.KSP().create(comm = comm)
    #create a 1d DMDA
    da = PETSc.DMDA().create(dim = 1,
                             boundary_type=(PETSc.DMDA.BoundaryType.NONE,),
                             sizes = (3,), dof = 1, stencil_width = 1,
                             comm = comm
                             )
    ksp.setDM(da)
    b = da.createGlobalVec()
    x = b.duplicate()
    #set the solver functions
    ksp.setComputeRHS(computeRHS,Ctx)
    ksp.setComputeOperators(computeMatrix,Ctx)
    ksp.setFromOptions()
    #solve
    ksp.solve(b,x)

    A,_ = ksp.getOperators()
    x = ksp.getSolution()
    b = ksp.getRhs()
    b2 = b.duplicate()
    A.mult(x,b2)
    b2.axpy(-1.0,x)
    norm = b2.norm(PETSc.NormType.NORM_MAX)
    PETSc.Sys.Print("Residual norm", norm)
    return 0

if __name__=="__main__":

    OptDB = PETSc.Options()#get PETSc option DB
    try:
        k = OptDB.getInt('k') #K coefficient
    except KeyError:
        k = OptDB.getInt('k',1)
    try:
        e = OptDB.getScalar('e') #e coefficient
    except KeyError:
        e = OptDB.getScalar('e',.99)

    main(k,e)

[Lisandro Dalcin]

You need to pass args,kargs as tuple/dict respectively. In your case, pass args=(Ctx,), as shown below:

ksp.setComputeRHS(computeRHS, args=(Ctx,))
ksp.setComputeOperators(computeMatrix, args=(Ctx,))

I agree this looks a bit unpythonic, however a more pythonic way to write your code would be:

class AppCtx():
    def __init__(self, k, e):
        self.k, self.e = k, e
    def computeMatrix(self, ksp, J, P):
        Ctx = self
        ...
    def computeRHS(self, ksp, b):
        Ctx = self
        ...

def main(k,e):
    ...
    Ctx = AppCtx(k,e)
    ...
    ksp.setComputeRHS(Ctx.computeRHS)
    ksp.setComputeOperators(Ctx.computeMatrix)
    ...

[Francesco Caimmi]

Ah, ok, it looks like I completely misunderstood what the Python interface does and missed the fact that here "args" is actually a ... keyword argument! Now if I run it with -ksp_view and -mat_view I get the same results for the C and Python versions with respect to the solver output and the the matrix entries. So I guess the exposed routines work as intended.

If you wish I could translate something slightly more complicated or with more degrees of freedom (such as ex28) to check if everything works as expected.

[Lisandro Dalcin]

Yes, please! Also, please consider contributing your example to the demo/ directory!

[Francesco Caimmi]

Of course it took forever, but finally I was able to find the time to translate ksp/tutorials/ex28.c. Everything seems to be ok, the results of the c code can be reproduced exactly. I would gladly contribute the two ksp examples to the demo directory:should I "simply" open a pull request? Something that should I know about best practices you adopt in doing so?

[Lisandro Dalcin]

I would say just create a directory demo/kspdm/, and your script there, commit and create a pull request.