We specialize matrix coloring and iterating in ArrayInterface.jl for structured matrices including:
DiagonalBidiagonalTridiagonalBandedMatrixBlockBandedMatrixBandedBlockBandedMatrix
The coloring function matrix_colors for structured matrices are overloaded by analytical results assuming all of the valid entries in the structured matrices are non-zero. Therefore, it is several magnitude faster than the default coloring algorithm. It is recommended to use the overloaded matrix_colors as long as the majority of the entries are filled by nonzero values. If not, you could further exploit the sparsity by convert the structured matrix into SparseCSC matrix and apply default coloring algorithm to it.
For diagonal (Diagonal, Bidiagonal and Tridiagonal) and banded (BandedMatrix) matrices, their color vectors have a cyclic form: [1,2,...,n,1,2,...,n,...] where n is the bandwidth of the matrix. The bandwidths of Diagonal, Bidiagonal and Tridiagonal is 1, 2 and 3 respectively, so their coloring vectors are [1,1,1,1,1,.....], [1,2,1,2,1,2,...] and [1,2,3,1,2,3,...]. The bandwidth of BandedMatrix is l+u+1 where l,u is lower and upper bandwidth of the matrix (l,u=BandedMatrices.bandwidths(A)). Note that this analytical coloring is not optimal for l<0 or u<0 as Jacobians with this shape could be easily reduced to smaller size.
Coloring for BlockBandedMatrix also has a cyclic form as it has a banded structure in a block level. For a BlockBandedMatrix with equal block widths w , the reptend would be [1,2,...,w,w+1,...,2*w,2*w+1,...,D*w] where D=lambda+mu+1 is the block bandwidth. For a BlockBandedMatrix with different block widths, we need to reserve colors for the maximum block widths of every D-th blocks. For example, the following BlockBandedMatrix has lower/upper block bandwidth of 0 and 1 with block widths of 1,2,3,4, so its color vector would be [1, 4,5, 1,2,3, 4,5,6,7].
+-+--+---+----+
|x|xx|ooo|oooo|
|x|xx|ooo|oooo|
|x|xx|ooo|oooo|
|x|xx|ooo|oooo|
+-+--+---+----+
|o|xx|xxx|oooo|
|o|xx|xxx|oooo|
|o|xx|xxx|oooo|
+-+--+---+----+
|o|oo|xxx|xxxx|
|o|oo|xxx|xxxx|
+-+--+---+----+
|o|oo|ooo|xxxx|
+-+--+---+----+
For a BandedBlockBandedMatrix, the only difference is that the colors inside a block could collapse from trivially increasing sequence to a cycling sequence due to its inner banded structure. The following BandedBlockBandedMatrix has lower/upper block bandwidth of 1 and 0, lower/upper subblock bandwidth of 1 and 1, and block widths of 4,3,2,1. Its analytical color vector would be [1,2,3,1, 4,5,6, 1,2, 4]
+----+---+--+-+
|xxoo|ooo|oo|o|
+----+---+--+-+
|xxoo|xxo|oo|o|
|xxxo|xxx|oo|o|
+----+---+--+-+
|oooo|xxo|xx|o|
|oooo|xxx|xx|o|
|oooo|oxx|ox|o|
+----+---+--+-+
|oooo|ooo|xx|x|
|oooo|ooo|xx|x|
|oooo|ooo|ox|o|
|oooo|ooo|oo|o|
+----+---+--+-+
A structured matrix A could be iterated like this:
rowinds,colinds=findstructralnz(A)
for i = 1:length(rowinds)
rowind=rowinds[i]
colind=colinds[i]
A[rowind,colind]=...
end
We implement findstructralnz(A) in ArrayInterface.jl for structured matrices. The function returns two indexable (implemented getindex) objects containing row and column indexes. The indexable object is a struct that stores the information needed to calculate the coordinates of the i-th entry. The getindex function (colinds[i] <=> getindex(colinds,i)) then calculates the coordinates on the fly. The speed of such iteration is dependent on two factors: cache missing rate and indexing speed. The first one could be optimized by giving entries in their memory order while the second one is fixed by the matrix type (and thus the type of indexable object). Therefore, this method is only efficient for those with fast indexing including Diagonal, Bidiagonal, Tridiagonal, BandedMatrix and SparseCSC. It would be very slow for BlockBandedMatrix and BandedBlockBandedMatrix.
Special treatment is made for BlockBandedMatrix and BandedBlockBandedMatrix in colored jacobian to keep iteration efficient. The code is coupled with colored jacobian so no general API is exposed. Generally, we exploit the block-shape memory layout to make it fast. Please check SciML/DifferentialEquations.jl#483 for further information.
In order to switch on colored jacobian when using OrdinaryDiffEq.jl, one has to supply an ODEFunction object with colorvec and jac_prototype specified. Usually the colorvec is an integer array indicating the color of each column of the Jacobian matrix while the jac_prototype is a matrix indicating the sparsity of the Jacobian matrix. If colorvec is specified, jac_prototype can not be a dense matrix since it does not hold sparsity information.