Skip to content

Automatic JuliaFormatter.jl run #67

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
github-actions wants to merge 1 commit into
base: main
Choose a base branch
from
+101 −70
Open

Automatic JuliaFormatter.jl run #67

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
164 changes: 98 additions & 66 deletions src/FastAlmostBandedMatrices.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,31 @@ module FastAlmostBandedMatrices

import PrecompileTools: @setup_workload, @compile_workload

using ArrayInterface, ArrayLayouts, BandedMatrices, ConcreteStructs, LazyArrays,
LinearAlgebra, MatrixFactorizations, Reexport

import ArrayLayouts: MemoryLayout, sublayout, sub_materialize, MatLdivVec, materialize!,
triangularlayout, triangulardata, zero!, _copyto!, colsupport,
rowsupport, _qr, _qr!, _factorize, muladd!
using ArrayInterface,
ArrayLayouts,
BandedMatrices,
ConcreteStructs,
LazyArrays,
LinearAlgebra,
MatrixFactorizations,
Reexport

import ArrayLayouts:
MemoryLayout,
sublayout,
sub_materialize,
MatLdivVec,
materialize!,
triangularlayout,
triangulardata,
zero!,
_copyto!,
colsupport,
rowsupport,
_qr,
_qr!,
_factorize,
muladd!
import BandedMatrices: _banded_qr!, bandeddata, banded_qr_lmul!
import LinearAlgebra: ldiv!
import MatrixFactorizations: QR, QRPackedQ, getQ, getR, QRPackedQLayout, AdjQRPackedQLayout
Expand All @@ -24,8 +43,8 @@ import MatrixFactorizations: QR, QRPackedQ, getQ, getR, QRPackedQLayout, AdjQRPa
A lazy representation of the union of two ranges, supporting iteration and indexing
without heap allocation.
"""
struct DisjointRange{T <: Integer, R1 <: AbstractUnitRange{T}, R2 <: AbstractUnitRange{T}} <:
AbstractVector{T}
struct DisjointRange{T<:Integer,R1<:AbstractUnitRange{T},R2<:AbstractUnitRange{T}} <:
AbstractVector{T}
r1::R1
r2::R2
end
Expand All @@ -39,7 +58,7 @@ Base.length(d::DisjointRange) = length(d.r1) + length(d.r2)
if i <= n1
return @inbounds d.r1[i]
else
return @inbounds d.r2[i - n1]
return @inbounds d.r2[i-n1]
end
end

Expand Down Expand Up @@ -120,14 +139,17 @@ where `2`'s are the fill part, and `1`'s are the bands part, and `3`'s are the o
part.
"""
@concrete struct AlmostBandedMatrix{T} <: LayoutMatrix{T}
bands
fill
bands::Any
fill::Any
end

function AlmostBandedMatrix(
::UndefInitializer, ::Type{T}, mn::NTuple{2, Integer},
lu::NTuple{2, Integer}, rank::Integer
) where {T}
::UndefInitializer,
::Type{T},
mn::NTuple{2,Integer},
lu::NTuple{2,Integer},
rank::Integer,
) where {T}
@assert lu[2] ≥ rank - 1
@assert rank ≥ 1 "Rank 0 fill array makes it a BandedMatrix."
bands = BandedMatrix{T}(undef, mn, lu)
Expand All @@ -136,17 +158,20 @@ function AlmostBandedMatrix(
end

function AlmostBandedMatrix{T}(
::UndefInitializer, mn::NTuple{2, Integer},
lu::NTuple{2, Integer}, rank::Integer
) where {T}
::UndefInitializer,
mn::NTuple{2,Integer},
lu::NTuple{2,Integer},
rank::Integer,
) where {T}
return AlmostBandedMatrix(undef, T, mn, lu, rank)
end

function AlmostBandedMatrix(
::UndefInitializer, mn::NTuple{2, Integer}, lu::NTuple{
2, Integer,
}, rank::Integer
)
::UndefInitializer,
mn::NTuple{2,Integer},
lu::NTuple{2,Integer},
rank::Integer,
)
return AlmostBandedMatrix(undef, Float64, mn, lu, rank)
end

Expand Down Expand Up @@ -178,7 +203,7 @@ end
@inline function finish_part_setindex!(bands, fill)
# copy `fill` into `bands` in the correct locations
l, u = bandwidths(bands)
for i in 1:size(fill, 1), j in max(1, i - l):min(size(bands, 2), i + u)
for i = 1:size(fill, 1), j = max(1, i-l):min(size(bands, 2), i+u)

@inbounds bands[i, j] = fill[i, j]
end
Expand Down Expand Up @@ -229,7 +254,7 @@ E = exclusive_bandpart(A) # Returns a view of rows 3:10 of the banded part
"""
@inline function exclusive_bandpart(A)
B, F = bandpart(A), fillpart(A)
return @view(B[(size(F, 1) + 1):end, :])
return @view(B[(size(F, 1)+1):end, :])
end

@inline almostbandwidths(_, A) = bandwidths(bandpart(A))
Expand Down Expand Up @@ -274,9 +299,9 @@ end
@inline function rowsupport(::AbstractAlmostBandedLayout, A, k)
l, _ = almostbandwidths(A)
if maximum(k) ≤ almostbandedrank(A)
return max(1, minimum(k) - l):size(A, 2)
return max(1, minimum(k)-l):size(A, 2)
else
return max(1, minimum(k) - l):min(maximum(k) + l, size(A, 2))
return max(1, minimum(k)-l):min(maximum(k)+l, size(A, 2))
end
end

Expand Down Expand Up @@ -316,12 +341,9 @@ end

# TODO: Support views properly
function sublayout(
::AlmostBandedLayout, ::Type{
<:Tuple{
AbstractUnitRange{Int}, AbstractUnitRange{Int},
},
}
)
::AlmostBandedLayout,
::Type{<:Tuple{AbstractUnitRange{Int},AbstractUnitRange{Int}}},
)
return AlmostBandedLayout()
end

Expand Down Expand Up @@ -357,9 +379,10 @@ end
# ---------------
function _almost_banded_summary(io, B::AlmostBandedMatrix{T}, inds) where {T}
return print(
io, Base.dims2string(length.(inds)),
io,
Base.dims2string(length.(inds)),
" AlmostBandedMatrix{$T} with bandwidths $(almostbandwidths(B)) and fill \
rank $(almostbandedrank(B))"
rank $(almostbandedrank(B))",
)
end
function Base.array_summary(io::IO, B::AlmostBandedMatrix, inds::Tuple{Vararg{Base.OneTo}})
Expand All @@ -384,7 +407,7 @@ end

function ArrayInterface.fast_scalar_indexing(A::AlmostBandedMatrix)
return ArrayInterface.fast_scalar_indexing(typeof(A.bands)) &&
ArrayInterface.fast_scalar_indexing(typeof(A.fill))
ArrayInterface.fast_scalar_indexing(typeof(A.fill))
end

function ArrayInterface.qr_instance(A::AlmostBandedMatrix{T}, pivot = NoPivot()) where {T}
Expand All @@ -406,7 +429,8 @@ function _almostbanded_qr(_, A)
# Expand the bandsize for the QR factorization
## Bypass the safety checks in `AlmostBandedMatrix`
return almostbanded_qr!(
AlmostBandedMatrix{eltype(A)}(BandedMatrix(copy(B), (l, l + u)), copy(L)), Val(true)
AlmostBandedMatrix{eltype(A)}(BandedMatrix(copy(B), (l, l + u)), copy(L)),
Val(true),
)
end

Expand Down Expand Up @@ -446,10 +470,10 @@ end

k = 1
while k ≤ ncols
kr = k:min(k + l + u, m)
jr1 = k:min(k + u, n)
jr2 = (k + u + 1):min(last(kr) + u, n)
jr3 = k:min(k + u, n, ncols)
kr = k:min(k+l+u, m)
jr1 = k:min(k+u, n)
jr2 = (k+u+1):min(last(kr)+u, n)
jr3 = k:min(k+u, n, ncols)
S = B[kr, jr1]
τv = τ[jr3]
R, _ = _banded_qr!(S, τv, length(jr3))
Expand All @@ -458,28 +482,24 @@ end
B_right = B[kr, jr2]
L_right = L[:, jr2]
U′ = U[kr, :]
for j in 1:length(jr2)
muladd!(
-one(T), U′[(j + 1):end, :], L_right[:, j], one(T), B_right[(j + 1):end, j]
)
for j = 1:length(jr2)
muladd!(-one(T), U′[(j+1):end, :], L_right[:, j], one(T), B_right[(j+1):end, j])
end
banded_qr_lmul!(Q', B_right)
banded_qr_lmul!(Q', U′)
for j in 1:length(jr2)
muladd!(
one(T), U′[(j + 1):end, :], L_right[:, j], one(T), B_right[(j + 1):end, j]
)
for j = 1:length(jr2)
muladd!(one(T), U′[(j+1):end, :], L_right[:, j], one(T), B_right[(j+1):end, j])
end
k = last(jr1) + 1
end

return AlmostBandedMatrix{eltype(A)}(B, Mul(U, L)), τ
end

function getQ(F::QR{<:Any, <:AlmostBandedMatrix})
function getQ(F::QR{<:Any,<:AlmostBandedMatrix})
return LinearAlgebra.QRPackedQ(bandpart(F.factors), F.τ)
end
function getR(F::QR{<:Any, <:AlmostBandedMatrix})
function getR(F::QR{<:Any,<:AlmostBandedMatrix})
n = min(size(F.factors, 1), size(F.factors, 2))
return UpperTriangular(view(F.factors, 1:n, 1:n))
end
Expand All @@ -502,11 +522,11 @@ end

_almostbanded_widerect_ldiv!(::QR{T}, B) where {T} = error("Not implemented")

const UpperLayoutMatrix{T} = UpperTriangular{T, <:LayoutMatrix{T}}
const UpperLayoutMatrix{T} = UpperTriangular{T,<:LayoutMatrix{T}}

for Typ in
(:StridedVector, :StridedMatrix, :AbstractVecOrMat, :UpperLayoutMatrix, :LayoutMatrix)
@eval function ldiv!(A::QR{T, <:AlmostBandedMatrix}, B::$Typ{T}) where {T}
@eval function ldiv!(A::QR{T,<:AlmostBandedMatrix}, B::$Typ{T}) where {T}
m, n = size(A)
return if m == n
_almostbanded_square_ldiv!(A, B)
Expand All @@ -527,7 +547,7 @@ function Base.materialize!(M::Lmul{<:AdjQRPackedQLayout{<:AlmostBandedLayout}})
return lmul!(QRPackedQ(bandpart(Q.factors), Q.τ)', M.B)
end

triangularlayout(::Type{Tri}, ::ML) where {Tri, ML <: AlmostBandedLayout} = Tri{ML}()
triangularlayout(::Type{Tri}, ::ML) where {Tri,ML<:AlmostBandedLayout} = Tri{ML}()

@inline function __arguments(x::LazyArray, ::AlmostBandedMatrix, ::Val)
return LazyArrays.arguments(x)
Expand All @@ -552,8 +572,11 @@ end
@inline __original_almostbandedrank(A) = size(first(__lowrankfillpart(A)), 2)

@views function _almostbanded_upper_ldiv!(
::Type{Tri}, R::AbstractMatrix, b::AbstractVector{T}, buffer
) where {T, Tri}
::Type{Tri},
R::AbstractMatrix,
b::AbstractVector{T},
buffer,
) where {T,Tri}
B = bandpart(R)
U, V = __lowrankfillpart(R)
fill!(buffer, zero(T))
Expand All @@ -562,9 +585,9 @@ end
k = n = size(R, 2)

while k > 0
kr = max(1, k - u):k
jr1 = (k + 1):(k + u + 1)
jr2 = (k + u + 2):(k + 2u + 2)
kr = max(1, k-u):k
jr1 = (k+1):(k+u+1)
jr2 = (k+u+2):(k+2u+2)
bv = b[kr]
if jr2[1] < n
muladd!(one(T), V[:, jr2], b[jr2], one(T), buffer)
Expand All @@ -580,30 +603,30 @@ end
return b
end

function Base.materialize!(M::MatLdivVec{TriangularLayout{'U', 'N', AlmostBandedLayout}})
function Base.materialize!(M::MatLdivVec{TriangularLayout{'U','N',AlmostBandedLayout}})
R, x = M.A, M.B
A = triangulardata(R)
r = __original_almostbandedrank(A)
_almostbanded_upper_ldiv!(UpperTriangular, A, x, Vector{eltype(M)}(undef, r))
return x
end

function Base.materialize!(M::MatLdivVec{TriangularLayout{'U', 'U', AlmostBandedLayout}})
function Base.materialize!(M::MatLdivVec{TriangularLayout{'U','U',AlmostBandedLayout}})
R, x = M.A, M.B
A = triangulardata(R)
r = __original_almostbandedrank(A)
_almostbanded_upper_ldiv!(UnitUpperTriangular, A, x, Vector{eltype(M)}(undef, r))
return x
end

function Base.materialize!(M::MatLdivVec{TriangularLayout{'L', 'N', AlmostBandedLayout}})
function Base.materialize!(M::MatLdivVec{TriangularLayout{'L','N',AlmostBandedLayout}})
R, x = M.A, M.B
A = triangulardata(R)
materialize!(Ldiv(LowerTriangular(bandpart(A)), x))
return x
end

function Base.materialize!(M::MatLdivVec{TriangularLayout{'L', 'U', AlmostBandedLayout}})
function Base.materialize!(M::MatLdivVec{TriangularLayout{'L','U',AlmostBandedLayout}})
R, x = M.A, M.B
A = triangulardata(R)
materialize!(Ldiv(UnitLowerTriangular(bandpart(A)), x))
Expand All @@ -615,11 +638,15 @@ end
# ---------------

@views function muladd!(
α, A::AlmostBandedMatrix, B::AbstractVecOrMat, β, C::AbstractVecOrMat
)
α,
A::AlmostBandedMatrix,
B::AbstractVecOrMat,
β,
C::AbstractVecOrMat,
)
L = fillpart(A)
muladd!(α, L, B, β, selectdim(C, 1, 1:size(L, 1)))
muladd!(α, exclusive_bandpart(A), B, β, selectdim(C, 1, (size(L, 1) + 1):size(C, 1)))
muladd!(α, exclusive_bandpart(A), B, β, selectdim(C, 1, (size(L, 1)+1):size(C, 1)))
return C
end

Expand All @@ -646,7 +673,12 @@ end
end
end

export AlmostBandedMatrix, bandpart, fillpart, exclusive_bandpart, finish_part_setindex!,
almostbandwidths, almostbandedrank
export AlmostBandedMatrix,
bandpart,
fillpart,
exclusive_bandpart,
finish_part_setindex!,
almostbandwidths,
almostbandedrank

end
7 changes: 3 additions & 4 deletions test/core_tests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,8 @@ end
B, L = bandpart(A), fillpart(A)

F = qr(A)
@test F.Q isa LinearAlgebra.QRPackedQ{Float64, <:BandedMatrix}
@test F.R isa UpperTriangular{Float64, <:SubArray{Float64, 2, <:AlmostBandedMatrix}}
@test F.Q isa LinearAlgebra.QRPackedQ{Float64,<:BandedMatrix}
@test F.R isa UpperTriangular{Float64,<:SubArray{Float64,2,<:AlmostBandedMatrix}}
@test F.Q' * A ≈ F.R
@test A == Ã

Expand All @@ -82,8 +82,7 @@ end
n = 80
A = AlmostBandedMatrix(BandedMatrix(fill(2.0, n, n), (1, 1)), fill(3.0, 1, n))
b = randn(n)
@test MemoryLayout(UpperTriangular(A)) ==
TriangularLayout{'U', 'N', AlmostBandedLayout}()
@test MemoryLayout(UpperTriangular(A)) == TriangularLayout{'U','N',AlmostBandedLayout}()
@test_broken UpperTriangular(Matrix(A)) \ b ≈ UpperTriangular(A) \ b
@test_broken UnitUpperTriangular(Matrix(A)) \ b ≈ UnitUpperTriangular(A) \ b
@test LowerTriangular(Matrix(A)) \ b ≈ LowerTriangular(A) \ b
Expand Down