<html><head></head><body><div class="ydp108a8da2yahoo-style-wrap" style="font-family:Helvetica Neue, Helvetica, Arial, sans-serif;font-size:13px;"><div></div>
        <div dir="ltr" data-setdir="false">Hi All,</div><div dir="ltr" data-setdir="false"><br></div><div dir="ltr" data-setdir="false">Thank you for the recommendations. The first option provided by Matthew works in petsc4py but required some additional computations to re-organize information. I developed an approach that was easier to build within my current system using the option provided by Hong along with some other posts I found on the user forum. See below for an example of how I integrated Numba and petcs4py so the filling of matrix elements on each processor is done using the optimized machine code produced by Numba (you can eliminate the cleaning step if the number of non-zero elements for each row is well defined). Scaling and overall performance is now satisfactory. <br></div><div dir="ltr" data-setdir="false"><br></div><div dir="ltr" data-setdir="false">I do have another question. In order to make this work I create the Petsc matrix 'A' twice, first to get local ownership and second to define the local elements for a particular processor:<br></div><div dir="ltr" data-setdir="false"><br></div><div dir="ltr" data-setdir="false">Algorithm</div><div dir="ltr" data-setdir="false">------------<br></div><div dir="ltr" data-setdir="false">(1) Create a Petsc matrix 'A' and set size and type</div><div dir="ltr" data-setdir="false">(2) Get the local row start and end for matrix 'A'</div><div dir="ltr" data-setdir="false">(3) Define the local non-zero coefficients for the rows owned by processor using a Numba JIT-compiled loop and store result in a csr matrix defined using Scipy. <br></div><div dir="ltr" data-setdir="false">(4) Redefine Petsc matix 'A' using the the local csr elements define in step 3.</div><div dir="ltr" data-setdir="false">(5) Begin and end assembly</div><div dir="ltr" data-setdir="false">(6) Define RHS and initial solution vectors</div><div dir="ltr" data-setdir="false">(7) Solve system <br></div><div dir="ltr" data-setdir="false"><span> (see code example below for details)</span></div><div dir="ltr" data-setdir="false"><span><br></span></div><div dir="ltr" data-setdir="false"><span>Steps (1) and (4) appear redundant and potentially sub-optimal from a performance perspective (perhaps due to my limited experience with petscy4py). Is there a better approach in terms of elegance and performance? Also, if this is the optimal syntax could someone elaborate on what exactly is occurring behind the seen in terms of memory allocation?<br></span></div><div dir="ltr" data-setdir="false"><span><br></span></div><div dir="ltr" data-setdir="false"><span>Thank you all again for your guidance and insights.<br></span></div><div dir="ltr" data-setdir="false"><span><br></span></div><div dir="ltr" data-setdir="false"><span><span><span><span>Modified Version of Dalcin's 2D Poisson Example</span></span></span> with Numba<br></span></div><div dir="ltr" data-setdir="false">----------------------------------------------------------------------------------<br></div><div dir="ltr" data-setdir="false"><div dir="ltr" data-setdir="false">import sys<br>from typing import Tuple<br>import numpy.typing as npt<br>import numpy as np<br>from numba import njit<br>import petsc4py<br>petsc4py.init(sys.argv)<br>from petsc4py import PETSc<br>import scipy.sparse as sps<br><br><br>def main() -> None:<br>    xnodes = 8000<br>    nnz_max = 10<br>    <br>    ynodes = xnodes<br>    N = xnodes*ynodes<br>    dx = 1.0/(xnodes + 1)<br>    <br>    A = PETSc.Mat()<br>    A.create(comm=PETSc.COMM_WORLD)<br>    A.setSizes((xnodes*ynodes, xnodes*ynodes))<br>    A.setType(PETSc.Mat.Type.AIJ)<br>    <br>    rstart, rend = A.getOwnershipRange()<br> <div>    Ascipy = build_csr_matrix(<br>        rstart, rend, nnz_max, dx, xnodes, ynodes)<br>    csr=(<br>        Ascipy.indptr[rstart:rend+1] - Ascipy.indptr[rstart],<br>        Ascipy.indices[Ascipy.indptr[rstart]:Ascipy.indptr[rend]],<br>        Ascipy.data[Ascipy.indptr[rstart]:Ascipy.indptr[rend]]<br>        )<br>    A = PETSc.Mat().createAIJ(size=(N,N), csr=csr)</div>    <br>    A.assemblyBegin()<br>    A.assemblyEnd()<br>    <br>    ksp = PETSc.KSP()<br>    ksp.create(comm=A.getComm())<br>    ksp.setType(PETSc.KSP.Type.CG)<br>    ksp.getPC().setType(PETSc.PC.Type.GAMG)<br>    <br>    ksp.setOperators(A)<br>    ksp.setFromOptions()<br>    <br>    x, b = A.createVecs()<br>    b.set(1.0)<br>    <br>    ksp.solve(b, x)<br>    <br>    <br>def build_csr_matrix(<br>        rstart:int, <br>        rend:int, <br>        nnz_max:int, <br>        dx:float, <br>        xnodes:int, <br>        ynodes:int<br>):<br>    Anz, Arow, Acol = build_nonzero_arrays(<br>        rstart, rend, nnz_max, dx, xnodes, ynodes<br>        )<br>    N = xnodes*ynodes<br>    Ls = sps.csr_matrix((Anz, (Arow,Acol)), shape=(N,N), dtype=np.float64)<br>    return Ls<br>  <br>  <br>def build_nonzero_arrays(<br>        rstart:int, <br>        rend:int, <br>        nnz_max:int, <br>        dx:float, <br>        xnodes:int, <br>        ynodes:int<br>) -> Tuple[<br>    npt.NDArray[np.float64], <br>    npt.NDArray[np.float64], <br>    npt.NDArray[np.float64]<br>    ]:<br>    nrows_local = (rend - rstart) + 1<br>    Anz_ini = np.zeros((nnz_max*nrows_local), dtype=np.float64)<br>    Arow_ini = np.zeros((nnz_max*nrows_local), dtype=np.int32)<br>    Acol_ini = np.zeros((nnz_max*nrows_local), dtype=np.int32)<br>    icount_nz = define_nonzero_values(<br>        rstart, rend, dx, xnodes, ynodes, Anz_ini, Arow_ini, Acol_ini<br>        )<br>    (<br>        Anz, Arow, Acol<br>    ) = clean_nonzero_arrays(icount_nz, Anz_ini, Arow_ini, Acol_ini)<br>    return Anz, Arow, Acol<br><br><br>@njit<br>def define_nonzero_values(<br>        rstart:int, <br>        rend:int,<br>        dx:float,<br>        xnodes:int,<br>        ynodes:int,<br>        Anz:npt.NDArray[np.float64],<br>        Arow:npt.NDArray[np.int64],<br>        Acol:npt.NDArray[np.int64]<br>) -> int:<br>    """ Fill matrix A<br>    """<br>    icount_nz = 0<br>    for row in range(rstart, rend):<br>        i, j = index_to_grid(row, xnodes)<br>        #A[row, row] = 4.0/dx**2<br>        Anz[icount_nz] = 4.0/dx**2<br>        Arow[icount_nz] = row<br>        Acol[icount_nz] = row<br>        icount_nz += 1<br>        if i > 0:<br>            column = row - xnodes<br>            #A[row, column] = -1.0/dx**2<br>            Anz[icount_nz] = -1.0/dx**2<br>            Arow[icount_nz] = row<br>            Acol[icount_nz] = column<br>            icount_nz += 1<br>        if i < xnodes - 1:<br>            column = row + xnodes<br>            #A[row, column] = -1.0/dx**2<br>            Anz[icount_nz] = -1.0/dx**2<br>            Arow[icount_nz] = row<br>            Acol[icount_nz] = column<br>            icount_nz += 1<br>        if j > 0:<br>            column = row - 1<br>            #A[row, column] = -1.0/dx**2<br>            Anz[icount_nz] = -1.0/dx**2<br>            Arow[icount_nz] = row<br>            Acol[icount_nz] = column<br>            icount_nz += 1<br>        if j < xnodes - 1:<br>            column = row + 1<br>            #A[row, column] = -1.0/dx**2<br>            Anz[icount_nz] = -1.0/dx**2<br>            Arow[icount_nz] = row<br>            Acol[icount_nz] = column<br>            icount_nz += 1<br>    return icount_nz<br><br><br>@njit<br>def clean_nonzero_arrays(<br>        icount_nz:int, <br>        Anz_ini:npt.NDArray[np.float64], <br>        Arow_ini:npt.NDArray[np.float64], <br>        Acol_ini:npt.NDArray[np.float64]<br>) -> Tuple[<br>    npt.NDArray[np.float64], <br>    npt.NDArray[np.float64], <br>    npt.NDArray[np.float64]<br>    ]:<br>    Anz = np.zeros((icount_nz), dtype=np.float64)<br>    Arow = np.zeros((icount_nz), dtype=np.int32)<br>    Acol = np.zeros((icount_nz), dtype=np.int32)<br>    for i in range(icount_nz):<br>        Anz[i] = Anz_ini[i]<br>        Arow[i] = Arow_ini[i]<br>        Acol[i] = Acol_ini[i]<br>    return Anz, Arow, Acol<br><br><br>@njit<br>def index_to_grid(r:int, n:int) -> Tuple[int,int]:<br>    """Convert a row number into a grid point."""<br>    return (r//n, r%n)<br><br><br>if __name__ == "__main__":<br>    main()</div><div><br></div></div><div><br></div>
        
        </div><div id="yahoo_quoted_2889210779" class="yahoo_quoted">
            <div style="font-family:'Helvetica Neue', Helvetica, Arial, sans-serif;font-size:13px;color:#26282a;">
                
                <div>
                    On Friday, August 18, 2023 at 11:24:36 AM CDT, Zhang, Hong <hongzhang@anl.gov> wrote:
                </div>
                <div><br></div>
                <div><br></div>
                <div><div dir="ltr">You can use this to build a PETSc matrix with the index arrays ai,aj and the value array aa:<br clear="none">    PETSc.Mat().createAIJ(size=(nrows,ncols), csr=(ai,aj,aa))<br clear="none"><br clear="none">Hong (Mr.)<br clear="none"><div class="yqt4861462648" id="yqtfd33436"><br clear="none">> On Aug 17, 2023, at 7:37 AM, Erik Kneller via petsc-users <<a shape="rect" ymailto="mailto:petsc-users@mcs.anl.gov" href="mailto:petsc-users@mcs.anl.gov">petsc-users@mcs.anl.gov</a>> wrote:<br clear="none">> <br clear="none">> Hi All,<br clear="none">> <br clear="none">> I need to fill non-zero values of a Petsc matrix via petsc4py for the domain defined by A.getOwnershipRange() using three Numpy arrays: (1) array containing row indices of non-zero value, (2) array containing column indices of non-zero values and (3) array containing the non-zero matrix values. How can one perform this type of filling operation in petsc4py? The method A.setValues does not appear to allow this since it only works on an individual matrix element or a block of matrix elements.<br clear="none">> <br clear="none">> I am using Numpy arrays since they can be computed in loops optimized using Numba on each processor. I also cannot pass the Petsc matrix to a Numba compiled function since type information cannot be inferred. I absolutely need to avoid looping in standard Python to define Petsc matrix elements due to performance issues. I also need to use a standard petscy4py method and avoid writing new C or Fortran wrappers to minimize language complexity.<br clear="none">> <br clear="none">> Example Algorithm Building on Lisandro Dalcin's 2D Poisson Example:<br clear="none">> ----------------------------------------------------------------------------------------------<br clear="none">> comm = PETSc.COMM_WORLD<br clear="none">> rank = comm.getRank()<br clear="none">> <br clear="none">> dx = 1.0/(xnodes + 1) # xnodes is the number of nodes in the x and y-directions of the grid<br clear="none">> nnz_max = 5 # max number of non-zero values per row<br clear="none">> <br clear="none">> A = PETSc.Mat()<br clear="none">> A.create(comm=PETSc.COMM_WORLD)<br clear="none">> A.setSizes((xnodes*ynodes, xnodes*ynodes))<br clear="none">> A.setType(PETSc.Mat.Type.AIJ)<br clear="none">> A.setPreallocationNNZ(nnz_max)<br clear="none">> <br clear="none">> rstart, rend = A.getOwnershipRange()<br clear="none">> <br clear="none">> # Here Anz, Arow and Acol are vectors with size equal to the number of non-zero values <br clear="none">> Anz, Arow, Acol = build_nonzero_numpy_arrays_using_numba(rstart, rend, nnz_max, dx, xnodes, ynodes)<br clear="none">> <br clear="none">> A.setValues(Arow, Acol, Anz) # <--- This does not work.<br clear="none">> <br clear="none">> A.assemblyBegin()<br clear="none">> A.assemblyEnd()<br clear="none">> ksp = PETSc.KSP()<br clear="none">> ksp.create(comm=A.getComm())<br clear="none">> ksp.setType(PETSc.KSP.Type.CG)<br clear="none">> ksp.getPC().setType(PETSc.PC.Type.GAMG)<br clear="none">> ksp.setOperators(A)<br clear="none">> ksp.setFromOptions()<br clear="none">> x, b = A.createVecs()<br clear="none">> b.set(1.0)<br clear="none">> <br clear="none">> ksp.solve(b, x)<br clear="none">> <br clear="none">> Regards,<br clear="none">> Erik Kneller, Ph.D.<br clear="none">> <br clear="none"><br clear="none"></div></div></div>
            </div>
        </div></body></html>