Creating and Solving Block Sparse Matrices with PETSc

General Usage
,
3

Thank you for the suggestion! I looked into KrylovKit,jl a bit, but haven’t implemented it in my full code. However, I did just implement it in the toy problem included in my post and while there is no error, the performance difference is so large that I feel I am likely implementing something incorrectly. I can’t properly benchmark PETSc due to the mentioned error but it using @time on successful runs I get the following results for solving the 400x400 block diagonal matrix 1000 times:

PETSc: ~0.25 S
KrylovKit: ~ 40 S

I would expect KrylovKit to perform slower since it doesn’t support preconditioning, but not 100x slower. At scale I also find that the manipulation of the sparse array I feed KrylovKit and IterativeSolvers is more expensive than the time needed to change the values of the PETSc matrix. The KrylovKit version of the toy code is below if you are interested:

using BenchmarkTools
using FLoops
using PETSc
using SparseArrays
using KrylovKit
using IncompleteLU

function SparseAssignment!(As, rowRange, colRange, M)
# Modifies sparse array using dense matrix ranges
# ASSUMES PROVIDED RANGES ARE FILLED ENTRIES IN PROVIDED SPARSE MATRIX
# Assigns each element per row per column

    # Loop over column range
    for j in eachindex(colRange)

        # Get indices of values in the current column
        ind_col = As.colptr[colRange[j]]:(As.colptr[colRange[j]+1]-1)

        # Get row number of each value in column
        rows = As.rowval[ind_col]

        # Loop through each row and assign values
        for i in eachindex(rowRange)

            nz_ind = ind_col[findfirst(rows.==rowRange[i])]

            As.nzval[nz_ind] = M[i,j]
        end

    end

end # Function SparseAssignment!

function PopulateArray!(A, nEq, nBlocks)
# Adjusts A entries along diagonal blocks of sparse array

    for i in 1:nBlocks
        
        # Get indices of block in array
        block_ind = ((i-1)*nEq + 1):(i*nEq)

        # Assign random values to sub-matrix
        SparseAssignment!(A, block_ind, block_ind, randn(Float64, nEq, nEq))

    end

end

function SolveArray(A, nEq, nBlocks)
    
    # Adjust values of A matrix
    PopulateArray!(A, nEq, nBlocks)

    # Predoncition system
    #LU = ilu(A; τ=0.1)

    # Create right-hand-side vector (Ax = rhs)
    rhs = randn(Float64, nEq*nBlocks)

    # Solve system
    x = linsolve(A, rhs)
    
end

function Test()
# Creates a block-diagonal system of equations and solves with KrylovKit
# -------------------------------------------------------------------

    # Number of solver iterations to run
    nit = 1000

    # System of equation parameters
    nEq = 4 # Size of sub-matrices
    nBlocks = 100 # Number of sub-matrices

    # Initialize block-diagonal sparse array
    AI = Int64[] # I indices
    AJ = Int64[] # J Indices
    # get I,J indices of each block along the diagonal
    for i in 1:nBlocks 
       rowRange = (nEq*(i-1)+1):(nEq*i)

       for ii in rowRange
           for jj in rowRange
           push!(AI, ii)
           push!(AJ, jj)
           end
       end
    end
    # Create sparse initialized with zeros
    A = sparse(AI, AJ, zeros(Float64, size(AI)))

    println("Starting Iterations")
    for i in 1:nit

        # Pass array to solver
        sol = SolveArray(A, nEq, nBlocks)

    end
    println("Iterations Concluded")

end

Test()