Skip to content

Commit

Permalink
Move block/banded stuff to requires
Browse files Browse the repository at this point in the history
  • Loading branch information
@Tokazama
Tokazama committed Jan 7, 2022
1 parent b5a73e0 commit 452903c
Show file tree
Hide file tree
Showing 2 changed files with 208 additions and 202 deletions.
204 changes: 204 additions & 0 deletions src/ArrayInterface.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -805,6 +805,53 @@ function __init__()
end

@require BandedMatrices = "aae01518-5342-5314-be14-df237901396f" begin
struct BandedMatrixIndex <: MatrixIndex
count::Int
rowsize::Int
colsize::Int
bandinds::Array{Int,1}
bandsizes::Array{Int,1}
isrow::Bool
end

Base.firstindex(i::BandedMatrixIndex) = 1
Base.lastindex(i::BandedMatrixIndex) = i.count
Base.length(i::BandedMatrixIndex) = lastindex(i)
@propagate_inbounds function Base.getindex(ind::BandedMatrixIndex, i::Int)
@boundscheck 1 <= i <= ind.count || throw(BoundsError(ind, i))
_i = i
p = 1
while _i - ind.bandsizes[p] > 0
_i -= ind.bandsizes[p]
p += 1
end
bandind = ind.bandinds[p]
startfromone = !xor(ind.isrow, (bandind > 0))
if startfromone
return _i
else
return _i + abs(bandind)
end
end

function _bandsize(bandind, rowsize, colsize)
-(rowsize - 1) <= bandind <= colsize - 1 || throw(ErrorException("Invalid Bandind"))
if (bandind * (colsize - rowsize) > 0) & (abs(bandind) <= abs(colsize - rowsize))
return min(rowsize, colsize)
elseif bandind * (colsize - rowsize) <= 0
return min(rowsize, colsize) - abs(bandind)
else
return min(rowsize, colsize) - abs(bandind) + abs(colsize - rowsize)
end
end

function BandedMatrixIndex(rowsize, colsize, lowerbandwidth, upperbandwidth, isrow)
upperbandwidth > -lowerbandwidth || throw(ErrorException("Invalid Bandwidths"))
bandinds = upperbandwidth:-1:-lowerbandwidth
bandsizes = [_bandsize(band, rowsize, colsize) for band in bandinds]
BandedMatrixIndex(sum(bandsizes), rowsize, colsize, bandinds, bandsizes, isrow)
end

function findstructralnz(x::BandedMatrices.BandedMatrix)
l, u = BandedMatrices.bandwidths(x)
rowsize, colsize = Base.size(x)
Expand All @@ -827,6 +874,82 @@ function __init__()

@require BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" begin
@require BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" begin
struct BlockBandedMatrixIndex <: MatrixIndex
count::Int
refinds::Array{Int,1}
refcoords::Array{Int,1}# storing col or row inds at ref points
isrow::Bool
end
Base.firstindex(i::BlockBandedMatrixIndex) = 1
Base.lastindex(i::BlockBandedMatrixIndex) = i.count
Base.length(i::BlockBandedMatrixIndex) = lastindex(i)
function BlockBandedMatrixIndex(nrowblock, ncolblock, rowsizes, colsizes, l, u)
blockrowind = BandedMatrixIndex(nrowblock, ncolblock, l, u, true)
blockcolind = BandedMatrixIndex(nrowblock, ncolblock, l, u, false)
sortedinds = sort(
[(blockrowind[i], blockcolind[i]) for i = 1:length(blockrowind)],
by = x -> x[1],
)
sort!(sortedinds, by = x -> x[2], alg = InsertionSort)# stable sort keeps the second index in order
refinds = Array{Int,1}()
refrowcoords = Array{Int,1}()
refcolcoords = Array{Int,1}()
rowheights = pushfirst!(copy(rowsizes), 1)
cumsum!(rowheights, rowheights)
blockheight = 0
blockrow = 1
blockcol = 1
currenti = 1
lastrowind = sortedinds[1][1] - 1
lastcolind = sortedinds[1][2]
for ind in sortedinds
rowind, colind = ind
if colind == lastcolind
if rowind > lastrowind
blockheight += rowsizes[rowind]
end
else
for j = blockcol:blockcol+colsizes[lastcolind]-1
push!(refinds, currenti)
push!(refrowcoords, blockrow)
push!(refcolcoords, j)
currenti += blockheight
end
blockcol += colsizes[lastcolind]
blockrow = rowheights[rowind]
blockheight = rowsizes[rowind]
end
lastcolind = colind
lastrowind = rowind
end
for j = blockcol:blockcol+colsizes[lastcolind]-1
push!(refinds, currenti)
push!(refrowcoords, blockrow)
push!(refcolcoords, j)
currenti += blockheight
end
push!(refinds, currenti)# guard
push!(refrowcoords, -1)
push!(refcolcoords, -1)
rowindobj = BlockBandedMatrixIndex(currenti - 1, refinds, refrowcoords, true)
colindobj = BlockBandedMatrixIndex(currenti - 1, refinds, refcolcoords, false)
rowindobj, colindobj
end
@propagate_inbounds function Base.getindex(ind::BlockBandedMatrixIndex, i::Int)
@boundscheck 1 <= i <= ind.count || throw(BoundsError(ind, i))
p = 1
while i - ind.refinds[p] >= 0
p += 1
end
p -= 1
_i = i - ind.refinds[p]
if ind.isrow
return ind.refcoords[p] + _i
else
return ind.refcoords[p]
end
end

function findstructralnz(x::BlockBandedMatrices.BlockBandedMatrix)
l, u = BlockBandedMatrices.blockbandwidths(x)
nrowblock = BlockBandedMatrices.blocksize(x, 1)
Expand All @@ -842,6 +965,87 @@ function __init__()
u,
)
end
struct BandedBlockBandedMatrixIndex <: MatrixIndex
count::Int
refinds::Array{Int,1}
refcoords::Array{Int,1}# storing col or row inds at ref points
reflocalinds::Array{BandedMatrixIndex,1}
isrow::Bool
end
Base.firstindex(i::BandedBlockBandedMatrixIndex) = 1
Base.lastindex(i::BandedBlockBandedMatrixIndex) = i.count
Base.length(i::BandedBlockBandedMatrixIndex) = lastindex(i)
@propagate_inbounds function Base.getindex(ind::BandedBlockBandedMatrixIndex, i::Int)
@boundscheck 1 <= i <= ind.count || throw(BoundsError(ind, i))
p = 1
while i - ind.refinds[p] >= 0
p += 1
end
p -= 1
_i = i - ind.refinds[p] + 1
ind.reflocalinds[p][_i] + ind.refcoords[p] - 1
end


function BandedBlockBandedMatrixIndex(
nrowblock,
ncolblock,
rowsizes,
colsizes,
l,
u,
lambda,
mu,
)
blockrowind = BandedMatrixIndex(nrowblock, ncolblock, l, u, true)
blockcolind = BandedMatrixIndex(nrowblock, ncolblock, l, u, false)
sortedinds = sort(
[(blockrowind[i], blockcolind[i]) for i = 1:length(blockrowind)],
by = x -> x[1],
)
sort!(sortedinds, by = x -> x[2], alg = InsertionSort)# stable sort keeps the second index in order
rowheights = pushfirst!(copy(rowsizes), 1)
cumsum!(rowheights, rowheights)
colwidths = pushfirst!(copy(colsizes), 1)
cumsum!(colwidths, colwidths)
currenti = 1
refinds = Array{Int,1}()
refrowcoords = Array{Int,1}()
refcolcoords = Array{Int,1}()
reflocalrowinds = Array{BandedMatrixIndex,1}()
reflocalcolinds = Array{BandedMatrixIndex,1}()
for ind in sortedinds
rowind, colind = ind
localrowind =
BandedMatrixIndex(rowsizes[rowind], colsizes[colind], lambda, mu, true)
localcolind =
BandedMatrixIndex(rowsizes[rowind], colsizes[colind], lambda, mu, false)
push!(refinds, currenti)
push!(refrowcoords, rowheights[rowind])
push!(refcolcoords, colwidths[colind])
push!(reflocalrowinds, localrowind)
push!(reflocalcolinds, localcolind)
currenti += localrowind.count
end
push!(refinds, currenti)
push!(refrowcoords, -1)
push!(refcolcoords, -1)
rowindobj = BandedBlockBandedMatrixIndex(
currenti - 1,
refinds,
refrowcoords,
reflocalrowinds,
true,
)
colindobj = BandedBlockBandedMatrixIndex(
currenti - 1,
refinds,
refcolcoords,
reflocalcolinds,
false,
)
rowindobj, colindobj
end

function findstructralnz(x::BlockBandedMatrices.BandedBlockBandedMatrix)
l, u = BlockBandedMatrices.blockbandwidths(x)
Expand Down
Loading

0 comments on commit 452903c

Please sign in to comment.